Searched refs:device_shape (Results 1 – 12 of 12) sorted by relevance
| /external/tensorflow/tensorflow/stream_executor/tpu/ |
| D | tpu_executable_interface.cc | 63 Shape device_shape = HostShapeToDeviceShape(host_shape); in AllocateOutputMemoryWithInputReuse() local
86 if (!Shape::Equal().MinorToMajorOnlyInLayout()(host_shape, device_shape)) { in AllocateOutputMemoryWithInputReuse()
89 << ShapeUtil::HumanStringWithLayout(device_shape); in AllocateOutputMemoryWithInputReuse()
93 ExecutionOutput result(std::move(device_shape), allocator, device_ordinal); in AllocateOutputMemoryWithInputReuse()
|
| D | tpu_executable.cc | 112 Shape device_shape = ApiConverter::FromC(&c_device_shape); in HostShapeToDeviceShape() local
115 return device_shape; in HostShapeToDeviceShape()
|
| D | tpu_transfer_manager.cc | 61 xla::Shape device_shape = ApiConverter::FromC(&c_device_shape); in HostShapeToDeviceShape() local
64 return device_shape; in HostShapeToDeviceShape()
|
| D | tpu_executor_c_api.h | 181 XLA_Shape* device_shape);
|
| /external/tensorflow/tensorflow/compiler/xla/service/ |
| D | transfer_manager.cc | 204 Shape* device_shape) { in ReadDynamicShapes() argument
205 DCHECK(device_shape->is_dynamic()); in ReadDynamicShapes()
206 Shape original_device_shape = *device_shape; in ReadDynamicShapes()
214 ShapeUtil::GetSubshape(*device_shape, index); in ReadDynamicShapes()
219 *ShapeUtil::GetMutableSubshape(device_shape, index); in ReadDynamicShapes()
248 device_shape->clear_dynamic_dimensions(); in ReadDynamicShapes()
250 TF_RET_CHECK(ShapeUtil::DynamicShapeIsCompatible(*device_shape, in ReadDynamicShapes()
|
| D | transfer_manager.h | 198 Shape* device_shape);
|
| /external/tensorflow/tensorflow/compiler/xrt/ |
| D | xrt_state.cc | 416 xla::Shape* device_shape) { in ExpandTreeOfTuples() argument
420 *device_shape = elements.shape(); in ExpandTreeOfTuples()
440 *xla::ShapeUtil::GetMutableSubshape(device_shape, index) = in ExpandTreeOfTuples()
463 xla::Shape device_shape; in MakeTuple() local
465 &host_shape, &device_shape)); in MakeTuple()
471 device_ordinal, allocator, host_shape, device_shape); in MakeTuple()
486 xla::ShapeUtil::GetSubshape(device_shape, index); in MakeTuple()
|
| D | xrt_state.h | 270 xla::Shape* device_shape);
|
| /external/tensorflow/tensorflow/core/tpu/kernels/ |
| D | tpu_execute_op.cc | 162 explicit InputBuffers(xla::Shape device_shape) in InputBuffers()
163 : buffers(std::move(device_shape)) {} in InputBuffers()
459 const xla::Shape& device_shape = in AllocateOutputTensors() local
467 xla::ScopedShapedBuffer shaped_buffer(device_shape, allocator, in AllocateOutputTensors()
|
| /external/tensorflow/tensorflow/core/tpu/ |
| D | tpu_execute.cc | 112 xla::Shape device_shape = ApiConverter::FromC(&c_device_shape); in HostShapeToDeviceShape() local
115 return device_shape; in HostShapeToDeviceShape()
|
| D | tpu_ops_c_api.h | 178 XLA_Shape* host_shape, XLA_Shape* device_shape);
|
| /external/tensorflow/tensorflow/compiler/jit/ |
| D | xla_launch_util.cc | 236 const xla::Shape& device_shape = in PopulateInputs() local
268 arguments.emplace_back(device_shape, shape); in PopulateInputs()
270 if (xla::Shape::Equal().MinorToMajorOnlyInLayout()(shape, device_shape)) { in PopulateInputs()
|