Searched refs:pjrt_client (Results 1 – 9 of 9) sorted by relevance
| /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() function
|
| D | BUILD | 133 "//tensorflow/compiler/xla/pjrt:pjrt_client",
205 "//tensorflow/compiler/xla/pjrt:pjrt_client",
235 "//tensorflow/compiler/xla/pjrt:pjrt_client",
266 "//tensorflow/compiler/xla/pjrt:pjrt_client",
317 "//tensorflow/compiler/xla/pjrt:pjrt_client",
364 "//tensorflow/compiler/xla/pjrt:pjrt_client",
389 "//tensorflow/compiler/xla/pjrt:pjrt_client",
533 "//tensorflow/compiler/xla/pjrt:pjrt_client",
|
| D | outfeed_receiver_py.cc | 58 return client->pjrt_client(); in OutfeedReceiverForPython()
101 return client->pjrt_client() == device->client(); in Callback()
|
| D | py_client.cc | 33 PyClient::PyClient(std::unique_ptr<PjRtClient> pjrt_client) in PyClient() argument
34 : pjrt_client_(std::move(pjrt_client)) {} in PyClient()
35 PyClient::PyClient(std::shared_ptr<PjRtClient> pjrt_client) in PyClient() argument
36 : pjrt_client_(std::move(pjrt_client)) {} in PyClient()
|
| D | jax_jit.cc | 573 h, pyclient.pjrt_client(), to_device), in HandleBool()
581 obj, pyclient.pjrt_client(), to_device), in HandleInt()
585 obj, pyclient.pjrt_client(), to_device), in HandleInt()
597 h, pyclient.pjrt_client(), to_device), in HandleFloat()
601 h, pyclient.pjrt_client(), to_device), in HandleFloat()
623 ValueOrThrow(pyclient.pjrt_client()->BufferFromHostBuffer( in HandleComplex()
632 ValueOrThrow(pyclient.pjrt_client()->BufferFromHostBuffer( in HandleComplex()
678 h, pyclient.pjrt_client(), to_device), in HandleNpBool()
682 h, pyclient.pjrt_client(), to_device), in HandleNpBool()
|
| D | dlpack.cc | 334 DeviceForDLContext(*client->pjrt_client(), dlmt->dl_tensor.ctx)); in DLPackManagedTensorToBuffer()
359 client->pjrt_client()->CreateViewOfDeviceBuffer( in DLPackManagedTensorToBuffer()
|
| D | xla_compiler.cc | 366 client->pjrt_client()->GetHloCostAnalysis()); in BuildXlaCompilerSubmodule()
|
| /external/tensorflow/tensorflow/compiler/xla/pjrt/ |
| D | BUILD | 132 name = "pjrt_client",
133 srcs = ["pjrt_client.cc"],
134 hdrs = ["pjrt_client.h"],
181 ":pjrt_client",
338 ":pjrt_client",
|
| /external/tensorflow/tensorflow/compiler/xla/python/tpu_driver/client/ |
| D | BUILD | 30 "//tensorflow/compiler/xla/pjrt:pjrt_client",
|