Home
last modified time | relevance | path

Searched refs:PyClient (Results 1 – 13 of 13) sorted by relevance

/external/tensorflow/tensorflow/compiler/xla/python/
Dxla.cc203 py::class_<PyClient, std::shared_ptr<PyClient>> py_local_client(m, "Client"); in PYBIND11_MODULE()
204 py_local_client.def_property_readonly("platform", &PyClient::platform_name) in PYBIND11_MODULE()
205 .def("device_count", &PyClient::device_count) in PYBIND11_MODULE()
206 .def("local_device_count", &PyClient::addressable_device_count) in PYBIND11_MODULE()
207 .def("devices", &PyClient::Devices) in PYBIND11_MODULE()
208 .def("local_devices", &PyClient::LocalDevices) in PYBIND11_MODULE()
209 .def("live_buffers", &PyClient::LiveBuffers) in PYBIND11_MODULE()
210 .def("host_id", &PyClient::task_id) in PYBIND11_MODULE()
211 .def("task_id", &PyClient::task_id) in PYBIND11_MODULE()
213 &PyClient::GetDefaultDeviceAssignment) in PYBIND11_MODULE()
[all …]
Dpy_client.h32 class PyClient; variable
69 std::shared_ptr<PyClient> client;
80 ClientAndPtr<T> WrapWithClient(std::shared_ptr<PyClient> client, T* contents) { in WrapWithClient()
89 class PyClient : public std::enable_shared_from_this<PyClient> {
91 explicit PyClient(std::unique_ptr<PjRtClient> pjrt_client);
92 explicit PyClient(std::shared_ptr<PjRtClient> pjrt_client);
Dpy_client.cc33 PyClient::PyClient(std::unique_ptr<PjRtClient> pjrt_client) in PyClient() function in xla::PyClient
35 PyClient::PyClient(std::shared_ptr<PjRtClient> pjrt_client) in PyClient() function in xla::PyClient
38 std::vector<ClientAndPtr<PjRtDevice>> PyClient::Devices() { in Devices()
48 std::vector<ClientAndPtr<PjRtDevice>> PyClient::LocalDevices() { in LocalDevices()
57 std::vector<ClientAndPtr<PyBuffer>> PyClient::LiveBuffers() { in LiveBuffers()
69 PyClient::GetDefaultDeviceAssignment(int num_replicas, int num_partitions) { in GetDefaultDeviceAssignment()
88 PyClient::GetDefaultDeviceAssignment1D(int num_replicas) { in GetDefaultDeviceAssignment1D()
102 StatusOr<std::unique_ptr<PjRtBuffer>> PyClient::PjRtBufferFromPyval( in PjRtBufferFromPyval()
146 StatusOr<std::unique_ptr<PyBuffer>> PyClient::BufferFromPyval( in BufferFromPyval()
158 StatusOr<std::shared_ptr<PyExecutable>> PyClient::Compile( in Compile()
[all …]
Dpy_executable.h38 PyExecutable(std::shared_ptr<PyClient> client,
44 std::shared_ptr<PyClient> client() const { return client_; } in client()
83 friend class PyClient;
85 std::shared_ptr<PyClient> client_;
Doutfeed_receiver_py.cc46 std::vector<std::shared_ptr<PyClient>> clients, in OutfeedReceiverForPython()
57 [](const std::shared_ptr<PyClient>& client) { in OutfeedReceiverForPython()
100 clients_, [device](const std::shared_ptr<PyClient>& client) { in Callback()
117 std::vector<std::shared_ptr<PyClient>> clients_;
129 std::vector<std::shared_ptr<PyClient>> clients, in BuildOutfeedReceiverSubmodule()
Dpy_buffer.h55 PyBuffer(std::shared_ptr<PyClient> client, std::unique_ptr<PjRtBuffer> buffer,
59 std::shared_ptr<PyClient> client() const { return client_; } in client()
113 friend class PyClient;
120 std::shared_ptr<PyClient> client_;
Djax_jit.cc568 py::handle, xla::PjRtDevice*, bool, xla::PyClient&)>;
571 bool jax_enable_x64, xla::PyClient& pyclient) { in HandleBool()
578 bool jax_enable_x64, xla::PyClient& pyclient) { in HandleInt()
594 xla::PyClient& pyclient) { in HandleFloat()
610 xla::PyClient& pyclient) { in HandleComplex()
643 xla::PyClient& pyclient) { in HandleDeviceArray()
665 xla::PyClient& pyclient) { in HandleBufferFromPyval()
675 bool jax_enable_x64, xla::PyClient& pyclient) { in HandleNpBool()
688 bool jax_enable_x64, xla::PyClient& pyclient) { in HandleUint64()
713 xla::PyClient& pyclient) { in HandleNdarray()
[all …]
Ddlpack.h32 const pybind11::capsule& tensor, std::shared_ptr<PyClient> client);
Djax_jit.h171 xla::PyClient& pyclient);
Dpy_executable.cc25 PyExecutable::PyExecutable(std::shared_ptr<PyClient> client, in PyExecutable()
Ddlpack.cc319 const pybind11::capsule& tensor, std::shared_ptr<PyClient> client) { in DLPackManagedTensorToBuffer()
Dpy_buffer.cc30 PyBuffer::PyBuffer(std::shared_ptr<PyClient> client, in PyBuffer()
Dxla_compiler.cc363 [](PyClient* client, in BuildXlaCompilerSubmodule()