Searched refs:tpu_input_shapes (Results 1 – 2 of 2) sorted by relevance
| /aosp_15_r20/external/tensorflow/tensorflow/core/tpu/kernels/ |
| H A D | tpu_functional_ops.cc | 402 std::map<int, InferredShape>& arg_shapes, EdgeShapes& tpu_input_shapes, in GetInputOutputInfo() argument
430 std::vector<int>& input_shapes = tpu_input_shapes[tpu_input_edge]; in GetInputOutputInfo()
652 const EdgeShapes& tpu_input_shapes, in GroupTensorsForInputPacking() argument
656 for (const auto& iter : tpu_input_shapes) { in GroupTensorsForInputPacking()
674 tpu_input_shapes.find(edge)->second; in GroupTensorsForInputPacking()
698 auto tpu_input_shape = tpu_input_shapes.find(edge)->second; in GroupTensorsForInputPacking()
746 Graph* graph, const string& cluster_name, EdgeShapes* tpu_input_shapes, in CreateConcatAndSplitNodesForInputTensor() argument
769 rank = tpu_input_shapes->at(edge).size(); in CreateConcatAndSplitNodesForInputTensor()
770 last_dim_vec.push_back(tpu_input_shapes->at(edge).back()); in CreateConcatAndSplitNodesForInputTensor()
899 (*tpu_input_shapes)[concat_to_split].push_back( in CreateConcatAndSplitNodesForInputTensor()
[all …]
|
| H A D | tpu_functional_ops.h | 56 const EdgeShapes& tpu_input_shapes,
64 Graph* graph, const string& cluster_name, EdgeShapes* tpu_input_shapes,
76 EdgeShapes* tpu_input_shapes,
|