ak.to_jax
---------

.. py:module: ak.to_jax

Defined in `awkward.operations.ak_to_jax <https://github.com/scikit-hep/awkward/blob/e9dd2694cf5c6ccb1717358fe926816d8ecbfb9c/src/awkward/operations/ak_to_jax.py>`__ on `line 13 <https://github.com/scikit-hep/awkward/blob/e9dd2694cf5c6ccb1717358fe926816d8ecbfb9c/src/awkward/operations/ak_to_jax.py#L13>`__.

.. py:function:: ak.to_jax(array)


    :param array: Array-like data (anything :py:obj:`ak.to_layout` recognizes).

Converts ``array`` (many types supported) into a JAX Device Array, if possible.

If the data are numerical and regular (nested lists have equal lengths
in each dimension, as described by the :py:obj:`ak.Array.type`), they can be losslessly
converted to a JAX array and this function returns without an error.

Otherwise, the function raises an error.

If ``array`` is a scalar, it is converted into a JAX scalar.

See also :py:obj:`ak.from_jax` and :py:obj:`ak.to_numpy`.