ak.to_jax — Awkward Array documentation

# ak.to_jax

Defined in awkward.operations.convert on line 595.

ak.to_jax(array)

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

If the data are numerical and regular (nested lists have equal lengths in each dimension, as described by the #type), they can be losslessly converted to a CuPy 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 `ak.to_cupy`, `ak.from_jax` and `ak.to_numpy`.