Searched refs:split_dim_node (Results 1 – 2 of 2) sorted by relevance
| /external/tensorflow/tensorflow/core/tpu/kernels/ |
| D | tpu_functional_ops.cc | 751 Node* split_dim_node = nullptr; in CreateConcatAndSplitNodesForInputTensor() local
757 .Finalize(graph, &split_dim_node)); in CreateConcatAndSplitNodesForInputTensor()
832 .Input(split_dim_node) in CreateConcatAndSplitNodesForInputTensor()
937 Node* split_dim_node = nullptr; in CreateConcatAndSplitNodesForOutputTensor() local
942 .Finalize(graph, &split_dim_node)); in CreateConcatAndSplitNodesForOutputTensor()
966 .Input(split_dim_node) in CreateConcatAndSplitNodesForOutputTensor()
1451 Node* split_dim_node = init_graph->AddNode(split_dim_def, &status); in InitializeShardedVarOnTPU() local
1452 split_dim_node->set_assigned_device_name(cpu_device); in InitializeShardedVarOnTPU()
1463 split_def.add_input(absl::StrCat(split_dim_node->name(), ":0")); in InitializeShardedVarOnTPU()
1469 init_graph->AddEdge(split_dim_node, 0, split_node, 0); in InitializeShardedVarOnTPU()
|
| /external/tensorflow/tensorflow/core/tpu/graph_rewrite/ |
| D | distributed_tpu_rewrite_pass.cc | 757 Node* split_dim_node = graph->AddNode(split_dim_def, &s); in CreateSplitNode() local
766 split_def.add_input(absl::StrCat(split_dim_node->name(), ":0")); in CreateSplitNode()
779 graph->AddEdge(split_dim_node, 0, split_node, 0); in CreateSplitNode()
786 graph->AddControlEdge(control_predecessor, split_dim_node); in CreateSplitNode()
|