| /external/tensorflow/tensorflow/compiler/xla/python/ |
| D | py_client.h | 91 explicit PyClient(std::unique_ptr<PjRtClient> pjrt_client);
92 explicit PyClient(std::shared_ptr<PjRtClient> pjrt_client);
94 PjRtClient* pjrt_client() const { return pjrt_client_.get(); } in pjrt_client()
95 std::shared_ptr<PjRtClient> shared_pjrt_client() { return pjrt_client_; } in shared_pjrt_client()
130 PjRtClient::HostBufferSemantics host_buffer_semantics);
133 PjRtClient::HostBufferSemantics host_buffer_semantics);
144 std::shared_ptr<PjRtClient> pjrt_client_;
|
| D | outfeed_receiver_test.cc | 35 PjRtClient* client) { in CompileAndExecute()
78 StatusOr<std::unique_ptr<PjRtClient>> GetCpuClientWithNonLocalDevice() { in GetCpuClientWithNonLocalDevice()
100 return std::unique_ptr<PjRtClient>(std::make_unique<PjRtStreamExecutorClient>( in GetCpuClientWithNonLocalDevice()
108 TF_ASSERT_OK_AND_ASSIGN(std::shared_ptr<PjRtClient> cpu_client, in TEST()
110 std::vector<PjRtClient*> clients{cpu_client.get()}; in TEST()
141 TF_ASSERT_OK_AND_ASSIGN(std::shared_ptr<PjRtClient> cpu_client, in TEST()
143 std::vector<PjRtClient*> clients{cpu_client.get()}; in TEST()
186 TF_ASSERT_OK_AND_ASSIGN(std::shared_ptr<PjRtClient> cpu_client, in TEST()
188 std::vector<PjRtClient*> clients{cpu_client.get()}; in TEST()
229 TF_ASSERT_OK_AND_ASSIGN(std::shared_ptr<PjRtClient> cpu_client, in TEST()
[all …]
|
| D | xla.cc | 196 py::enum_<PjRtClient::HostBufferSemantics>(m, "HostBufferSemantics") in PYBIND11_MODULE()
198 PjRtClient::HostBufferSemantics::kImmutableOnlyDuringCall) in PYBIND11_MODULE()
200 PjRtClient::HostBufferSemantics::kImmutableUntilTransferCompletes) in PYBIND11_MODULE()
201 .value("ZERO_COPY", PjRtClient::HostBufferSemantics::kZeroCopy); in PYBIND11_MODULE()
225 PjRtClient::HostBufferSemantics::kZeroCopy) in PYBIND11_MODULE()
233 TF_ASSIGN_OR_RETURN(std::unique_ptr<PjRtClient> client, in PYBIND11_MODULE()
239 TF_ASSIGN_OR_RETURN(std::unique_ptr<PjRtClient> client, in PYBIND11_MODULE()
249 std::unique_ptr<PjRtClient> client, in PYBIND11_MODULE()
260 TF_ASSIGN_OR_RETURN(std::shared_ptr<PjRtClient> client, in PYBIND11_MODULE()
|
| D | py_client.cc | 33 PyClient::PyClient(std::unique_ptr<PjRtClient> pjrt_client) in PyClient()
35 PyClient::PyClient(std::shared_ptr<PjRtClient> pjrt_client) in PyClient()
104 PjRtClient::HostBufferSemantics host_buffer_semantics) { in PjRtBufferFromPyval()
128 PjRtClient::HostBufferSemantics::kImmutableOnlyDuringCall) { in PjRtBufferFromPyval()
148 PjRtClient::HostBufferSemantics host_buffer_semantics) { in BufferFromPyval()
|
| D | outfeed_receiver.h | 47 OutfeedReceiver(Callback callback, absl::Span<PjRtClient* const> clients,
|
| D | outfeed_receiver.cc | 155 absl::Span<PjRtClient* const> clients,
229 OutfeedReceiver::Callback callback, absl::Span<PjRtClient* const> clients, in OutfeedReceiverImpl()
457 absl::Span<PjRtClient* const> clients, in OutfeedReceiver()
|
| D | jax_jit.cc | 359 const py::handle& scalar, xla::PjRtClient* client, in ConvertToScalarBuffer()
369 xla::PjRtClient::HostBufferSemantics::kImmutableOnlyDuringCall, nullptr, in ConvertToScalarBuffer()
625 xla::PjRtClient::HostBufferSemantics::kImmutableOnlyDuringCall, in HandleComplex()
634 xla::PjRtClient::HostBufferSemantics::kImmutableOnlyDuringCall, in HandleComplex()
670 xla::PjRtClient::HostBufferSemantics::kZeroCopy)); in HandleBufferFromPyval()
694 xla::PjRtClient::HostBufferSemantics::kZeroCopy)); in HandleUint64()
705 xla::PjRtClient::HostBufferSemantics::kZeroCopy)); in HandleUint64()
734 xla::PjRtClient::HostBufferSemantics::kZeroCopy)); in HandleNdarray()
|
| D | outfeed_receiver_py.cc | 55 std::vector<PjRtClient*> client_ptrs(clients_.size()); in OutfeedReceiverForPython()
|
| D | dlpack.cc | 225 StatusOr<PjRtDevice*> DeviceForDLContext(const PjRtClient& client, in DeviceForDLContext()
|
| /external/tensorflow/tensorflow/compiler/xla/pjrt/ |
| D | gpu_multistream_test.cc | 31 std::unique_ptr<PjRtClient> client, in TEST()
76 PjRtClient::HostBufferSemantics::kImmutableUntilTransferCompletes, in TEST()
82 PjRtClient::HostBufferSemantics::kImmutableUntilTransferCompletes, in TEST()
88 PjRtClient::HostBufferSemantics::kImmutableUntilTransferCompletes, in TEST()
|
| D | pjrt_client.h | 58 class PjRtClient; variable
65 virtual PjRtClient* client() const = 0;
139 class PjRtClient {
141 virtual ~PjRtClient() = default;
280 virtual PjRtClient* client() const = 0;
423 virtual PjRtClient* client() const = 0;
|
| D | interpreter_device.cc | 32 StatusOr<std::unique_ptr<PjRtClient>> GetInterpreterClient() { in GetInterpreterClient()
54 return std::unique_ptr<PjRtClient>(std::make_unique<PjRtStreamExecutorClient>( in GetInterpreterClient()
|
| D | cpu_device.cc | 32 StatusOr<std::unique_ptr<PjRtClient>> GetCpuClient(bool asynchronous) { in GetCpuClient()
60 return std::unique_ptr<PjRtClient>(std::make_unique<PjRtStreamExecutorClient>( in GetCpuClient()
|
| D | pjrt_stream_executor_client.h | 71 void SetClient(PjRtClient* client) { in SetClient()
84 PjRtClient* client() const override { return client_; } in client()
118 PjRtClient* client_ = nullptr;
121 class PjRtStreamExecutorClient : public PjRtClient {
461 PjRtClient* client, PjRtDevice* device);
545 friend class PjRtClient;
|
| D | interpreter_device.h | 32 StatusOr<std::unique_ptr<PjRtClient>> GetInterpreterClient();
|
| D | cpu_device.h | 31 StatusOr<std::unique_ptr<PjRtClient>> GetCpuClient(bool asynchronous);
|
| D | gpu_device.h | 57 StatusOr<std::unique_ptr<PjRtClient>> GetGpuClient(
|
| D | tpu_client.h | 54 StatusOr<std::shared_ptr<PjRtClient>> GetTpuClient(
|
| D | tpu_client.cc | 175 StatusOr<std::shared_ptr<PjRtClient>> GetTpuClient( in GetTpuClient()
219 return std::shared_ptr<PjRtClient>( in GetTpuClient()
|
| D | gpu_device.cc | 313 StatusOr<std::unique_ptr<PjRtClient>> GetGpuClient( in GetGpuClient()
336 return std::unique_ptr<PjRtClient>(std::make_unique<GpuClient>( in GetGpuClient()
|
| D | pjrt_stream_executor_client.cc | 366 bool is_uninitialized_create, PjRtClient* client, in AllocateDestinationBuffer()
999 PjRtClient* client, PjRtDevice* device) in PjRtStreamExecutorBuffer()
1481 PjRtClient* client, LocalDeviceState* local_device, in MakeTupleHelper()
1552 std::shared_ptr<BufferSequencingEvent> definition_event, PjRtClient* client, in OutputBufferHelper()
|
| /external/tensorflow/tensorflow/compiler/xla/python/tpu_driver/client/ |
| D | tpu_client.h | 56 PjRtClient* client() const override { return nullptr; } in client()
|