Searched defs:xla_module (Results 1 – 2 of 2) sorted by relevance
235 const auto xla_module = py::module::import("jax.interpreters.xla"); in HandleDeviceArray() local271 py::object xla_module = py::module::import("jax.interpreters.xla"); in DevicePut() local432 py::object xla_module = py::module::import("jax.interpreters.xla"); in PyArgSignatureOfValue() local
566 py::module xla_module(py::module::import("jax.interpreters.xla")); in ComputeSignature() local914 static const auto* xla_module = in Call() local