Searched refs:split_dim_node (Results 1 – 2 of 2) sorted by relevance
| /aosp_15_r20/external/tensorflow/tensorflow/core/tpu/kernels/ |
| H A D | tpu_functional_ops.cc | 801 Node* split_dim_node = nullptr; in CreateConcatAndSplitNodesForInputTensor() local
807 .Finalize(graph, &split_dim_node)); in CreateConcatAndSplitNodesForInputTensor()
882 .Input(split_dim_node) in CreateConcatAndSplitNodesForInputTensor()
987 Node* split_dim_node = nullptr; in CreateConcatAndSplitNodesForOutputTensor() local
992 .Finalize(graph, &split_dim_node)); in CreateConcatAndSplitNodesForOutputTensor()
1016 .Input(split_dim_node) in CreateConcatAndSplitNodesForOutputTensor()
1503 TF_ASSIGN_OR_RETURN(Node * split_dim_node, in InitializeShardedVarOnTPU()
1505 split_dim_node->set_assigned_device_name(cpu_device); in InitializeShardedVarOnTPU()
1515 split_def.add_input(absl::StrCat(split_dim_node->name(), ":0")); in InitializeShardedVarOnTPU()
1520 init_graph->AddEdge(split_dim_node, 0, split_node, 0); in InitializeShardedVarOnTPU()
|
| /aosp_15_r20/external/tensorflow/tensorflow/core/tpu/graph_rewrite/ |
| H A D | distributed_tpu_rewrite_pass.cc | 757 TF_ASSIGN_OR_RETURN(Node * split_dim_node, graph->AddNode(split_dim_def)); in CreateSplitNode()
765 split_def.add_input(absl::StrCat(split_dim_node->name(), ":0")); in CreateSplitNode()
777 graph->AddEdge(split_dim_node, 0, split_node, 0); in CreateSplitNode()
784 graph->AddControlEdge(control_predecessor, split_dim_node); in CreateSplitNode()
|