| /aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/service/ |
| H A D | allocation_tracker.cc | 35 ScopedShapedBuffer shaped_buffer, const std::string& tag) { in Register()
70 for (auto& shaped_buffer : replicated_buffers) { in RegisterInternal() local
103 for (const auto& shaped_buffer : replicated_buffers) { in Unregister() local
123 for (auto& shaped_buffer : it->second) { in Unregister() local
137 const ShapedBuffer* shaped_buffer = replicated_buffers[0]; in DeconstructTuple() local
196 for (const auto& shaped_buffer : it->second) { in ResolveInternal() local
|
| H A D | transfer_manager.cc | 167 ShapedBuffer shaped_buffer(on_device_shape, in TransferArrayToDeviceAsync() local
191 ShapedBuffer shaped_buffer(shape, stream->parent()->device_ordinal()); in TransferArrayFromDevice() local
394 ScopedShapedBuffer shaped_buffer(std::move(on_device_shape), allocator, in AllocateScopedShapedBuffer() local
|
| H A D | shaped_buffer.cc | 129 ScopedShapedBuffer::ScopedShapedBuffer(ShapedBuffer shaped_buffer, in ScopedShapedBuffer()
152 ShapedBuffer shaped_buffer(static_cast<ShapedBuffer&&>(*this)); in release() local
|
| H A D | executable.cc | 64 ShapedBuffer shaped_buffer(input_shape, device_ordinal); in ToShapedBuffer() local
93 const ShapedBuffer& shaped_buffer) { in MakeMaybeOwningDeviceMemoryTree()
|
| H A D | shaped_buffer_test.cc | 189 xla::ScopedShapedBuffer shaped_buffer(shape, /*allocator=*/&allocator, in BM_TakeSubTree() local
|
| H A D | service.cc | 253 const ShapedBuffer* shaped_buffer = replicated_buffers[replica]; in ResolveAndValidateArguments() local
|
| /aosp_15_r20/external/tensorflow/tensorflow/compiler/xrt/ |
| H A D | xrt_state.cc | 187 auto shaped_buffer = scoped_buffer->release(); in CreateAndTransfer() local
210 auto shaped_buffer = scoped_buffer->release(); in CreateUninitialized() local
221 const xla::ShapedBuffer& shaped_buffer, const xla::Shape& on_host_shape, in CreateFromBuffer()
234 const xla::ShapedBuffer& shaped_buffer, xla::Backend* backend, in CreateFromBuffer()
318 auto shaped_buffer = scoped_buffer->release(); in SwapIn() local
574 const xla::ShapedBuffer& shaped_buffer, in InitializeFromShapedBuffer()
587 xla::ShapedBuffer shaped_buffer(on_host_shape(), on_device_shape(), in ToShapedBuffer() local
|
| /aosp_15_r20/external/tensorflow/tensorflow/compiler/jit/ |
| H A D | xla_tensor.h | 61 const xla::ShapedBuffer& shaped_buffer() const { in shaped_buffer() function
65 xla::ShapedBuffer& shaped_buffer() { in shaped_buffer() function
70 void set_shaped_buffer(xla::ScopedShapedBuffer shaped_buffer) { in set_shaped_buffer()
|
| H A D | xla_tensor.cc | 50 xla::ScopedShapedBuffer shaped_buffer(on_host_shape, on_device_shape, in AllocateShapedBuffer() local
|
| H A D | xla_device.cc | 74 const xla::ShapedBuffer& shaped_buffer = xla_tensor->shaped_buffer(); in DefaultPaddedShapeFn() local
|
| /aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/pjrt/ |
| H A D | tracked_device_buffer.cc | 100 ScopedShapedBuffer* shaped_buffer, in FromScopedShapedBuffer()
124 ShapedBuffer shaped_buffer(on_device_shape, device_ordinal_); in AsShapedBuffer() local
|
| H A D | tfrt_cpu_pjrt_client.cc | 848 ShapedBuffer shaped_buffer(on_device_shape, device_ordinal); in AsShapedBuffer() local
880 ShapedBuffer shaped_buffer = AsShapedBuffer( in logical_on_device_shape() local
|
| H A D | pjrt_stream_executor_client.cc | 610 ShapedBuffer shaped_buffer = device_buffer->AsShapedBuffer(on_device_shape_); in logical_on_device_shape() local
1334 ShapedBuffer shaped_buffer = device_buffer->AsShapedBuffer(on_device_shape_); in ToLiteral() local
|
| /aosp_15_r20/external/tensorflow/tensorflow/core/tpu/kernels/ |
| H A D | tpu_execute_op.cc | 174 xla::ShapedBuffer shaped_buffer(std::move(host_shape), buffers.shape(), in ToShapedBuffer() local
466 xla::ScopedShapedBuffer shaped_buffer(device_shape, allocator, in AllocateOutputTensors() local
685 auto shaped_buffer = input_buffers->ToShapedBuffer(std::move(host_shape), in DoWork() local
|
| H A D | tpu_reshard_variables_op.cc | 187 xla::ShapedBuffer shaped_buffer(std::move(host_shape), input_buffers.shape(), in DoTpuExecute() local
|
| H A D | tpu_reshard_variables_op_util.cc | 282 xla::ScopedShapedBuffer shaped_buffer(host_shape, device_shape, allocator, in UpdateOutputVariables() local
|
| /aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/tests/ |
| H A D | local_client_test_base.cc | 144 const ShapedBuffer& shaped_buffer) { in ShapedBufferToLiteral()
|
| H A D | buffer_donation_test.cc | 107 ShapedBuffer shaped_buffer = scoped_shaped_buffer.release(); in RunAndCheck() local
|
| /aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/client/ |
| H A D | local_client.cc | 463 const ShapedBuffer& shaped_buffer) { in ShapedBufferToLiteral()
|
| /aosp_15_r20/external/tensorflow/tensorflow/compiler/xrt/kernels/ |
| H A D | xrt_execute_op.cc | 271 xla::ScopedShapedBuffer* shaped_buffer = run_result.MutableResult(); in CreateOutputTuple() local
|