Searched refs:device_buffers (Results 1 – 4 of 4) sorted by relevance
| /external/tensorflow/tensorflow/compiler/xla/pjrt/ |
| D | tracked_device_buffer_test.cc | 32 std::vector<stream_executor::DeviceMemoryBase> device_buffers; in MakeArray() local
42 device_buffers.push_back(device_memory.Release()); in MakeArray()
47 device_buffers, in MakeArray()
|
| D | pjrt_stream_executor_client.cc | 1483 absl::Span<const PjRtStreamExecutorBuffer::ScopedHold> device_buffers, in MakeTupleHelper() argument
1528 device_buffers) { in MakeTupleHelper()
1638 absl::Span<const PjRtStreamExecutorBuffer::ScopedHold> device_buffers, in MakeExecutionInputsAndWaitForEvents() argument
1648 device_buffers, device_ordinal)); in MakeExecutionInputsAndWaitForEvents()
1663 device_buffers[i].AddToInput( in MakeExecutionInputsAndWaitForEvents()
1685 std::vector<PjRtStreamExecutorBuffer::ScopedHold>* device_buffers, in EnqueueExecution() argument
1698 device_buffers->reserve(argument_handles.size()); in EnqueueExecution()
1709 device_buffers->emplace_back(handle->GetBufferWithHold( in EnqueueExecution()
1713 device_buffers->back(); in EnqueueExecution()
1746 device_ordinal, options, argument_handles, *device_buffers, events)); in EnqueueExecution()
[all …]
|
| D | pjrt_stream_executor_client.h | 685 absl::Span<const PjRtStreamExecutorBuffer::ScopedHold> device_buffers,
692 std::vector<PjRtStreamExecutorBuffer::ScopedHold>* device_buffers,
|
| /external/tensorflow/tensorflow/compiler/xla/python/ |
| D | pmap_lib.h | 163 pybind11::list device_buffers) in ShardedDeviceArray() argument
167 device_buffers_(device_buffers) {} in ShardedDeviceArray()
|