Searched refs:output_device_shape (Results 1 – 4 of 4) sorted by relevance
| /external/tensorflow/tensorflow/compiler/xrt/kernels/ |
| D | tpu_execute_op.cc | 144 xla::Shape output_device_shape = output_shaped_buffer.on_device_shape(); in AllocateOutputTuple() local
145 if (!output_device_shape.is_static()) { in AllocateOutputTuple()
148 stream, &output_shaped_buffer, &output_device_shape)); in AllocateOutputTuple()
153 xla::ShapeUtil::DeviceShapeToHostShape(output_device_shape); in AllocateOutputTuple()
156 output_shaped_buffer, output_host_shape, output_device_shape, in AllocateOutputTuple()
|
| D | xrt_execute_op.cc | 275 xla::Shape output_device_shape = shaped_buffer->on_device_shape(); in CreateOutputTuple() local
280 stream, shaped_buffer, &output_device_shape)); in CreateOutputTuple()
283 xla::ShapeUtil::DeviceShapeToHostShape(output_device_shape), in CreateOutputTuple()
284 output_device_shape, backend, device_ordinal, &output_tuple)); in CreateOutputTuple()
|
| /external/tensorflow/tensorflow/core/tpu/kernels/ |
| D | tpu_execute_op.cc | 442 xla::Shape output_device_shape = output_buffers->buffers.on_device_shape(); in AllocateOutputTensors() local
444 if (!output_device_shape.is_static()) { in AllocateOutputTensors()
446 stream, &output_buffers->buffers, &output_device_shape)); in AllocateOutputTensors()
449 xla::ShapeUtil::GetSubshape(output_device_shape, {i}); in AllocateOutputTensors()
460 xla::ShapeUtil::GetTupleElementShape(output_device_shape, i); in AllocateOutputTensors()
|
| /external/tensorflow/tensorflow/compiler/jit/ |
| D | xla_launch_util.cc | 468 xla::Shape output_device_shape = output.on_device_shape(); in PopulateOutputs() local
470 stream, &output, &output_device_shape)); in PopulateOutputs()
472 output.set_shapes(output_device_shape, output_device_shape); in PopulateOutputs()
475 xla::ShapeUtil::GetSubshape(output_device_shape, {i}); in PopulateOutputs()
|