ak.jax.import_jax
-----------------

.. py:module: ak.jax.import_jax

Defined in `awkward.jax <https://github.com/scikit-hep/awkward/blob/76cb66b946cbf2ae355bddf2530b323465395bae/src/awkward/jax.py>`__ on `line 135 <https://github.com/scikit-hep/awkward/blob/76cb66b946cbf2ae355bddf2530b323465395bae/src/awkward/jax.py#L135>`__.

.. py:function:: ak.jax.import_jax()

Ensure that JAX integration is registered, and return the JAX module. Raise a RuntimeError if not.