Searched refs:kXlaReplicaIdAttrName (Results 1 – 4 of 4) sorted by relevance
| /external/tensorflow/tensorflow/compiler/tf2xla/ |
| D | side_effect_util.h | 42 extern const char kXlaReplicaIdAttrName[];
|
| D | side_effect_util.cc | 29 const char kXlaReplicaIdAttrName[] = "_xla_replica_id"; variable
|
| /external/tensorflow/tensorflow/core/tpu/graph_rewrite/ |
| D | distributed_tpu_rewrite_pass.cc | 386 if (TryGetNodeAttr(node->def(), kXlaReplicaIdAttrName, &replica_id)) { in FindTaggedNodes()
3477 image->AddAttr(kXlaReplicaIdAttrName, replica_index); in CopyOutsideCompilationNodes()
3539 node->AddAttr(kXlaReplicaIdAttrName, 0); in ReplicateOutsideCompilationNodes()
3749 GetNodeAttr(n->def(), kXlaReplicaIdAttrName, &replica_id)); in LowerOutsideCompilationFunctionalNodes()
3818 node->AddAttr(kXlaReplicaIdAttrName, replica_id); in LowerOutsideCompilationFunctionalNodes()
3824 node->AddAttr(kXlaReplicaIdAttrName, replica_id); in LowerOutsideCompilationFunctionalNodes()
3850 GetNodeAttr(n->def(), kXlaReplicaIdAttrName, &replica_id)); in LowerOutsideCompilationFunctionalNodes()
4184 GetNodeAttr(node->def(), kXlaReplicaIdAttrName, &replica_id)); in UpdateHeadTailOutsideCompilation()
|
| D | encapsulate_tpu_computations_pass.cc | 474 copy_node->AddAttr(kXlaReplicaIdAttrName, replica_id); in MoveHeadOutsideCompilationToHost()
1125 copy_node->AddAttr(kXlaReplicaIdAttrName, replica_id); in MoveTailOutsideCompilationToHost()
|