Home
last modified time | relevance | path

Searched defs:on_device_shape (Results 1 – 15 of 15) sorted by relevance

/aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/service/
H A Dshaped_buffer.cc34 ShapedBuffer::ShapedBuffer(Shape on_device_shape, int device_ordinal) in ShapedBuffer()
41 ShapedBuffer::ShapedBuffer(Shape on_host_shape, Shape on_device_shape, in ShapedBuffer()
116 ScopedShapedBuffer::ScopedShapedBuffer(Shape on_device_shape, in ScopedShapedBuffer()
123 Shape on_device_shape, in ScopedShapedBuffer()
H A Dshaped_buffer.h68 const Shape& on_device_shape() const { return on_device_shape_; } in on_device_shape() function
100 void set_shapes(const Shape& on_device_shape) { in set_shapes()
109 void set_shapes(const Shape& on_host_shape, const Shape& on_device_shape) { in set_shapes()
H A Dexecutable.h155 ExecutionOutput(Shape on_host_shape, Shape on_device_shape, in ExecutionOutput()
158 ExecutionOutput(Shape on_device_shape, se::DeviceMemoryAllocator* allocator, in ExecutionOutput()
H A Dtransfer_manager.cc156 const Shape on_device_shape = HostShapeToDeviceShape(literal.shape()); in TransferArrayToDeviceAsync() local
389 Shape on_device_shape = (shape_representation_fn == nullptr) in AllocateScopedShapedBuffer() local
/aosp_15_r20/external/tensorflow/tensorflow/compiler/xrt/
H A Dxrt_state.cc88 xla::Shape on_device_shape = transfer_manager->HostShapeToDeviceShape(shape); in AllocateScopedShapedBuffer() local
151 const xla::Shape& on_device_shape) in XRTTupleAllocation()
222 const xla::Shape& on_device_shape, xla::Backend* backend, in CreateFromBuffer()
353 const xla::Shape& XRTTupleAllocation::on_device_shape() const { in on_device_shape() function in tensorflow::XRTTupleAllocation
/aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/pjrt/
H A Dtfrt_cpu_pjrt_client.cc432 const Shape& on_device_shape, in AllocateDestinationBuffer()
649 Shape on_device_shape, in TfrtCpuBuffer()
846 int device_ordinal, const Shape& on_device_shape, in AsShapedBuffer()
H A Dtfrt_cpu_pjrt_client.h293 const Shape& on_device_shape() const override { return on_device_shape_; } in on_device_shape() function
H A Dpjrt_stream_executor_client.h566 const Shape& on_device_shape() const override { return on_device_shape_; } in on_device_shape() function
H A Dpjrt_stream_executor_client.cc400 Shape on_device_shape = dst_buffer.on_device_shape(); in AllocateDestinationBuffer() local
1115 Shape on_device_shape, std::shared_ptr<TrackedDeviceBuffer> device_buffer, in PjRtStreamExecutorBuffer()
H A Dpjrt_c_api_client.cc633 const Shape& PjRtCApiBuffer::on_device_shape() const { in on_device_shape() function in xla::PjRtCApiBuffer
/aosp_15_r20/external/tensorflow/tensorflow/compiler/jit/
H A Dxla_tensor.cc45 const xla::Shape& on_device_shape, in AllocateShapedBuffer()
H A Dxla_tpu_device.cc90 const xla::Shape& on_device_shape = in TpuPaddedShapeFn() local
/aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/stream_executor/tpu/
H A Dtpu_executable_interface.cc159 const Shape& on_device_shape = result.Result().on_device_shape(); in AllocateOutputMemoryWithInputReuse() local
H A Dc_api_decl.h257 XLA_Shape on_device_shape; member
/aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/service/interpreter/
H A Dexecutable_base.cc213 const Shape& on_device_shape = result.Result().on_device_shape(); in AllocateOutputMemoryWithInputReuse() local