Searched refs:PjRtTpuDevice (Results 1 – 3 of 3) sorted by relevance
164 py::class_<PjRtTpuDevice, PjRtDevice, ClientAndPtr<PjRtTpuDevice>>( in PYBIND11_MODULE()166 .def_property_readonly("host_id", &PjRtTpuDevice::task_id) in PYBIND11_MODULE()167 .def_property_readonly("task_id", &PjRtTpuDevice::task_id) in PYBIND11_MODULE()170 [](const PjRtTpuDevice& device) -> pybind11::tuple { in PYBIND11_MODULE()175 "core_on_chip", &PjRtTpuDevice::core_on_chip, in PYBIND11_MODULE()177 .def("__repr__", [](const PjRtTpuDevice& device) { in PYBIND11_MODULE()
29 class PjRtTpuDevice : public PjRtStreamExecutorDevice {31 PjRtTpuDevice(const tensorflow::tpu::TpuCoreLocationExternal core, in PjRtTpuDevice() function
165 auto device = absl::make_unique<PjRtTpuDevice>( in GetTpuDevices()