Searched refs:xla_module (Results 1 – 1 of 1) sorted by relevance
387 const auto xla_module = py::module::import("jax.interpreters.xla"); in ArgSignatureOfValue() local388 const auto& device_array = xla_module.attr("_DeviceArray"); in ArgSignatureOfValue()751 const auto xla_module = py::module::import("jax.interpreters.xla"); in DevicePut() local752 const auto& device_array = xla_module.attr("_DeviceArray"); in DevicePut()952 static const auto* xla_module = in ConvertArgsToBuffers() local954 const auto& device_array = xla_module->attr("_DeviceArray"); in ConvertArgsToBuffers()1195 static const auto* xla_module = in Call() local1198 new py::handle(xla_module->attr("_DeviceArray")); in Call()