Searched refs:DeviceArrayBase (Results 1 – 6 of 6) sorted by relevance
42 class DeviceArrayBase {44 DeviceArrayBase() = default;53 class PyBuffer : public DeviceArrayBase {
155 class ShardedDeviceArray : xla::DeviceArrayBase {164 : DeviceArrayBase(), in ShardedDeviceArray()
266 py::class_<DeviceArrayBase> device_array_base(m, "DeviceArrayBase"); in PYBIND11_MODULE()269 py::class_<PyBuffer, DeviceArrayBase, std::unique_ptr<PyBuffer>> buffer( in PYBIND11_MODULE()
425 DeviceArrayBase = _xla.DeviceArrayBase variable
445 (*p)[py::type::handle_of<xla::DeviceArrayBase>().ptr()] = buffer_handler; in ArgSignatureOfValue()765 (*p)[py::type::handle_of<xla::DeviceArrayBase>().ptr()] = in DevicePut()
473 if not isinstance(buffer, xla_client.DeviceArrayBase):