Searched refs:output_tensor_shapes (Results 1 – 5 of 5) sorted by relevance
| /aosp_15_r20/external/tensorflow/tensorflow/core/tpu/kernels/ |
| H A D | tpu_execute_op.cc | 414 std::vector<TensorShape> output_tensor_shapes; in AllocateOutputTensors() local
415 output_tensor_shapes.reserve(sub_elements); in AllocateOutputTensors()
428 output_tensor_shapes.push_back(shape); in AllocateOutputTensors()
451 output_tensor_shapes[i] = shape; in AllocateOutputTensors()
485 TF_RET_CHECK(num_updated_variables <= output_tensor_shapes.size()) in AllocateOutputTensors()
486 << num_updated_variables << " <= " << output_tensor_shapes.size(); in AllocateOutputTensors()
520 for (int i = 0; i < output_tensor_shapes.size(); ++i) { in AllocateOutputTensors()
529 op_output_index, output_tensor_shapes[i], &output_tensor)); in AllocateOutputTensors()
554 output_tensor_shapes[i], in AllocateOutputTensors()
568 op_output_index, output_tensor_shapes[i], &output_tensor)); in AllocateOutputTensors()
[all …]
|
| H A D | tpu_reshard_variables_op_util.cc | 246 std::vector<TensorShape> output_tensor_shapes; in UpdateOutputVariables() local
247 output_tensor_shapes.reserve(sub_elements); in UpdateOutputVariables()
260 output_tensor_shapes.push_back(shape); in UpdateOutputVariables()
301 variables[i].var()->tensor()->dtype(), output_tensor_shapes[i], in UpdateOutputVariables()
|
| H A D | tpu_executable_info.proto | 53 repeated TensorShapeProto output_tensor_shapes = 8; field
|
| H A D | tpu_reshard_variables_op.cc | 269 context, std::move(result), executable->output_tensor_shapes(), backend, in DoTpuExecute()
|
| /aosp_15_r20/external/tensorflow/tensorflow/compiler/jit/ |
| H A D | xla_launch_util.cc | 505 std::vector<TensorShape> output_tensor_shapes; in PopulateOutputs() local
506 output_tensor_shapes.reserve(ctx->num_outputs()); in PopulateOutputs()
530 output_tensor_shapes.push_back(shape); in PopulateOutputs()
534 output_tensor_shapes.push_back(compilation_result->outputs[i].shape); in PopulateOutputs()
541 const TensorShape& shape = output_tensor_shapes[i]; in PopulateOutputs()
|