Searched refs:kTPUReplicateAttr (Results 1 – 7 of 7) sorted by relevance
| /external/tensorflow/tensorflow/compiler/mlir/tensorflow/transforms/ |
| D | tpu_cluster_cleanup_attributes.cc | 30 constexpr char kTPUReplicateAttr[] = "_tpu_replicate"; variable
41 for (StringRef attr : {kTPUReplicateAttr, kDeviceAttr}) in runOnOperation()
|
| D | tpu_cluster_formation.cc | 52 constexpr char kTPUReplicateAttr[] = "_tpu_replicate"; variable
94 auto tpu_replicate_attr = attrs.get(kTPUReplicateAttr); in CollectMetadata()
113 << kTPUReplicateAttr << "' attribute '" in CollectMetadata()
125 if (auto attr = op.getAttrOfType<StringAttr>(kTPUReplicateAttr)) { in CollectAndGroupClusterOps()
128 << "attribute '" << kTPUReplicateAttr << "' is empty"; in CollectAndGroupClusterOps()
275 inner_op->removeAttr(kTPUReplicateAttr); in CreateClusterOp()
547 << "TPUReplicateMetadata for associated '" << kTPUReplicateAttr in FormClustersInBlock()
|
| /external/tensorflow/tensorflow/compiler/tf2xla/ |
| D | mlir_bridge_pass.cc | 39 constexpr char kTPUReplicateAttr[] = "_tpu_replicate"; variable
54 op->getAttrOfType<mlir::StringAttr>(kTPUReplicateAttr); in HasTPUOp()
|
| /external/tensorflow/tensorflow/core/tpu/ |
| D | tpu_defs.cc | 27 const char* const kTPUReplicateAttr = "_tpu_replicate"; variable
|
| D | tpu_defs.h | 50 extern const char* const kTPUReplicateAttr;
|
| /external/tensorflow/tensorflow/core/tpu/graph_rewrite/ |
| D | encapsulate_tpu_computations_pass.cc | 241 AddNodeAttr(kTPUReplicateAttr, call_def->name(), call_def); in RewriteSubgraph()
475 copy_node->AddAttr(kTPUReplicateAttr, cluster_name); in MoveHeadOutsideCompilationToHost()
1126 copy_node->AddAttr(kTPUReplicateAttr, cluster_name); in MoveTailOutsideCompilationToHost()
1619 TF_RETURN_IF_ERROR(GetNodeAttr(n->def(), kTPUReplicateAttr, &cluster_name)); in RenameClustersWithDuplicatedNames()
1654 n->ClearAttr(kTPUReplicateAttr); in RenameClustersWithDuplicatedNames()
1655 n->AddAttr(kTPUReplicateAttr, new_cluster_name); in RenameClustersWithDuplicatedNames()
1659 if (GetNodeAttr(e->dst()->def(), kTPUReplicateAttr, &cluster_name) in RenameClustersWithDuplicatedNames()
2511 e->src()->attrs().Find(kTPUReplicateAttr) != nullptr && in Encapsulate()
2513 e->dst()->attrs().Find(kTPUReplicateAttr) == nullptr && in Encapsulate()
2533 kTPUReplicateAttr, **graph, RewriteSubgraph, in Encapsulate()
[all …]
|
| D | distributed_tpu_rewrite_pass.cc | 324 const AttrValue* cluster_attr = node->attrs().Find(kTPUReplicateAttr); in FindTaggedNodes()
327 kTPUReplicateAttr, " attr."); in FindTaggedNodes()
331 return errors::Internal("Attr ", kTPUReplicateAttr, " on node ", in FindTaggedNodes()
337 "TPUReplicate node ", node->name(), " has ", kTPUReplicateAttr, in FindTaggedNodes()
350 const AttrValue* cluster_attr = node->attrs().Find(kTPUReplicateAttr); in FindTaggedNodes()
357 kTPUReplicateAttr, " attr."); in FindTaggedNodes()
362 return errors::Internal("Attr ", kTPUReplicateAttr, " on node ", in FindTaggedNodes()
368 "Attr ", kTPUReplicateAttr, " on node ", node->name(), in FindTaggedNodes()
373 kTPUReplicateAttr, " attr but no ", in FindTaggedNodes()
3843 TF_RETURN_IF_ERROR(GetNodeAttr(n->def(), kTPUReplicateAttr, &replicate)); in LowerOutsideCompilationFunctionalNodes()
[all …]
|