Searched refs:xla_node (Results 1 – 1 of 1) sorted by relevance
| /external/tensorflow/tensorflow/core/tpu/graph_rewrite/ |
| D | encapsulate_tpu_computations_pass.cc | 333 Node* xla_node) { in UpdateMirroredVariableIndices() argument
335 if (xla_node->attrs().Find(TPUREPLICATE_MIRRORED_VAR_INDICES_ATTR) != in UpdateMirroredVariableIndices()
337 TF_RETURN_IF_ERROR(GetNodeAttr(xla_node->def(), in UpdateMirroredVariableIndices()
345 xla_node->ClearAttr(TPUREPLICATE_MIRRORED_VAR_INDICES_ATTR); in UpdateMirroredVariableIndices()
346 xla_node->AddAttr(TPUREPLICATE_MIRRORED_VAR_INDICES_ATTR, in UpdateMirroredVariableIndices()
359 const std::string& cluster_name, Graph* g, Graph* xla_graph, Node* xla_node, in MoveHeadOutsideCompilationToHost() argument
450 TF_RETURN_IF_ERROR(xla_node->input_edges(&input_edges)); in MoveHeadOutsideCompilationToHost()
453 TF_RETURN_IF_ERROR(GetNodeAttr(xla_node->attrs(), "Tinputs", &input_types)); in MoveHeadOutsideCompilationToHost()
455 TF_RETURN_IF_ERROR(GetNodeAttr(xla_node->attrs(), "num_distributed_variables", in MoveHeadOutsideCompilationToHost()
459 GetNodeAttr(xla_node->attrs(), "num_replicas", &num_replicas)); in MoveHeadOutsideCompilationToHost()
[all …]
|