Searched refs:pxla_module (Results 1 – 1 of 1) sorted by relevance
762 const auto pxla_module = py::module::import("jax.interpreters.pxla"); in DevicePut() local763 const auto& sda = pxla_module.attr("ShardedDeviceArray"); in DevicePut()