Searched refs:ClientAndPtr (Results 1 – 9 of 9) sorted by relevance
| /external/tensorflow/tensorflow/compiler/xla/python/ |
| D | py_client.h | 56 struct ClientAndPtr { struct
57 ClientAndPtr() = default;
60 explicit ClientAndPtr(T*) { in ClientAndPtr() function
64 ClientAndPtr(const ClientAndPtr&) = default;
65 ClientAndPtr(ClientAndPtr&&) = default;
66 ClientAndPtr& operator=(const ClientAndPtr&) = default; argument
67 ClientAndPtr& operator=(ClientAndPtr&&) = default; argument
80 ClientAndPtr<T> WrapWithClient(std::shared_ptr<PyClient> client, T* contents) { in WrapWithClient() argument
81 ClientAndPtr<T> result; in WrapWithClient()
106 std::vector<ClientAndPtr<PjRtDevice>> Devices();
[all …]
|
| D | py_client.cc | 38 std::vector<ClientAndPtr<PjRtDevice>> PyClient::Devices() { in Devices()
39 std::vector<ClientAndPtr<PjRtDevice>> devices; in Devices()
48 std::vector<ClientAndPtr<PjRtDevice>> PyClient::LocalDevices() { in LocalDevices()
49 std::vector<ClientAndPtr<PjRtDevice>> devices; in LocalDevices()
57 std::vector<ClientAndPtr<PyBuffer>> PyClient::LiveBuffers() { in LiveBuffers()
59 std::vector<ClientAndPtr<PyBuffer>> buffers; in LiveBuffers()
68 StatusOr<std::vector<std::vector<ClientAndPtr<PjRtDevice>>>>
73 std::vector<std::vector<ClientAndPtr<PjRtDevice>>> result; in GetDefaultDeviceAssignment()
87 StatusOr<std::vector<ClientAndPtr<PjRtDevice>>>
92 std::vector<ClientAndPtr<PjRtDevice>> result; in GetDefaultDeviceAssignment1D()
|
| D | xla.cc | 105 py::class_<PjRtDevice, ClientAndPtr<PjRtDevice>>( in PYBIND11_MODULE()
127 [](const ClientAndPtr<PjRtDevice>& device) { return device.client; }) in PYBIND11_MODULE()
154 py::class_<CpuDevice, PjRtDevice, ClientAndPtr<CpuDevice>>(m, "CpuDevice") in PYBIND11_MODULE()
159 py::class_<GpuDevice, PjRtDevice, ClientAndPtr<GpuDevice>>(m, "GpuDevice") in PYBIND11_MODULE()
164 py::class_<PjRtTpuDevice, PjRtDevice, ClientAndPtr<PjRtTpuDevice>>( in PYBIND11_MODULE()
|
| D | py_buffer.h | 62 ClientAndPtr<PjRtDevice> device() const;
69 const ClientAndPtr<PjRtDevice>& dst_device) const;
|
| D | py_executable.cc | 61 std::vector<ClientAndPtr<PjRtDevice>> PyExecutable::AddressableDevices() const { in AddressableDevices()
62 std::vector<ClientAndPtr<PjRtDevice>> devices; in AddressableDevices()
|
| D | py_executable.h | 51 std::vector<ClientAndPtr<PjRtDevice>> AddressableDevices() const;
|
| D | py_buffer.cc | 67 ClientAndPtr<PjRtDevice> PyBuffer::device() const { in device()
72 const ClientAndPtr<PjRtDevice>& dst_device) const { in CopyToDevice()
|
| D | outfeed_receiver_py.cc | 43 std::function<void(ClientAndPtr<PjRtDevice>, uint32_t, pybind11::object)>;
|
| D | jax_jit.cc | 911 xla::ClientAndPtr<xla::PjRtDevice> default_pydevice_;
1123 default_pydevice_ = py::cast<xla::ClientAndPtr<xla::PjRtDevice>>( in Call()
1236 xla::ClientAndPtr<xla::PjRtDevice> to_device) { in BuildJaxjitSubmodule()
|