Home
last modified time | relevance | path

Searched refs:argument_handles (Results 1 – 5 of 5) sorted by relevance

/external/tensorflow/tensorflow/compiler/xla/python/tpu_driver/client/
Dtpu_client.cc622 absl::Span<PyTpuBuffer* const> argument_handles) { in Execute() argument
639 PyTpuBuffer::MakeTuple(argument_handles, client_, in Execute()
643 all_core_arguments = std::vector<PyTpuBuffer*>(argument_handles.begin(), in Execute()
644 argument_handles.end()); in Execute()
647 ExecuteHelper(absl::MakeSpan(&all_core_arguments, 1), argument_handles, in Execute()
668 absl::Span<const std::vector<PyTpuBuffer*>> argument_handles) { in ExecuteOnLocalDevices() argument
674 if (argument_handles.size() != num_local_devices) { in ExecuteOnLocalDevices()
678 argument_handles.size(), num_local_devices, num_replicas(), in ExecuteOnLocalDevices()
689 tupled_arguments.resize(argument_handles.size()); in ExecuteOnLocalDevices()
690 tupled_argument_pointers.resize(argument_handles.size()); in ExecuteOnLocalDevices()
[all …]
Dtpu_client.h338 absl::Span<PyTpuBuffer* const> argument_handles);
346 absl::Span<const std::vector<PyTpuBuffer*>> argument_handles);
/external/tensorflow/tensorflow/compiler/xla/pjrt/
Dpjrt_stream_executor_client.cc1637 absl::Span<PjRtBuffer* const> argument_handles, in MakeExecutionInputsAndWaitForEvents() argument
1647 MakeTupleHelper(client_, device_state, argument_handles, in MakeExecutionInputsAndWaitForEvents()
1652 execution_inputs.reserve(argument_handles.size()); in MakeExecutionInputsAndWaitForEvents()
1653 for (int i = 0; i < argument_handles.size(); ++i) { in MakeExecutionInputsAndWaitForEvents()
1654 PjRtBuffer* handle = argument_handles[i]; in MakeExecutionInputsAndWaitForEvents()
1682 absl::Span<PjRtBuffer* const> argument_handles, int replica, int partition, in EnqueueExecution() argument
1698 device_buffers->reserve(argument_handles.size()); in EnqueueExecution()
1699 for (int i = 0; i < argument_handles.size(); ++i) { in EnqueueExecution()
1701 tensorflow::down_cast<PjRtStreamExecutorBuffer*>(argument_handles[i]); in EnqueueExecution()
1735 if (argument_handles.size() != 1) { in EnqueueExecution()
[all …]
Dpjrt_stream_executor_client.h651 absl::Span<const std::vector<PjRtBuffer*>> argument_handles,
655 absl::Span<PjRtBuffer* const> argument_handles, PjRtDevice* device,
659 absl::Span<PjRtBuffer* const> argument_handles, PjRtDevice* device,
684 absl::Span<PjRtBuffer* const> argument_handles,
689 absl::Span<PjRtBuffer* const> argument_handles, int replica,
702 absl::Span<PjRtBuffer* const> argument_handles, int replica,
Dpjrt_client.h461 Execute(absl::Span<const std::vector<PjRtBuffer*>> argument_handles,
468 absl::Span<PjRtBuffer* const> argument_handles, PjRtDevice* device,
475 absl::Span<PjRtBuffer* const> argument_handles, PjRtDevice* device,