/aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/service/ |
H A D | shaped_buffer.cc | 34 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 D | shaped_buffer.h | 68 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 D | executable.h | 155 ExecutionOutput(Shape on_host_shape, Shape on_device_shape, in ExecutionOutput() 158 ExecutionOutput(Shape on_device_shape, se::DeviceMemoryAllocator* allocator, in ExecutionOutput()
|
H A D | transfer_manager.cc | 156 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 D | xrt_state.cc | 88 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 D | tfrt_cpu_pjrt_client.cc | 432 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 D | tfrt_cpu_pjrt_client.h | 293 const Shape& on_device_shape() const override { return on_device_shape_; } in on_device_shape() function
|
H A D | pjrt_stream_executor_client.h | 566 const Shape& on_device_shape() const override { return on_device_shape_; } in on_device_shape() function
|
H A D | pjrt_stream_executor_client.cc | 400 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 D | pjrt_c_api_client.cc | 633 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 D | xla_tensor.cc | 45 const xla::Shape& on_device_shape, in AllocateShapedBuffer()
|
H A D | xla_tpu_device.cc | 90 const xla::Shape& on_device_shape = in TpuPaddedShapeFn() local
|
/aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/stream_executor/tpu/ |
H A D | tpu_executable_interface.cc | 159 const Shape& on_device_shape = result.Result().on_device_shape(); in AllocateOutputMemoryWithInputReuse() local
|
H A D | c_api_decl.h | 257 XLA_Shape on_device_shape; member
|
/aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/service/interpreter/ |
H A D | executable_base.cc | 213 const Shape& on_device_shape = result.Result().on_device_shape(); in AllocateOutputMemoryWithInputReuse() local
|