ak.to_jax

Defined in awkward.operations.ak_to_jax on line 13.

ak.to_jax(array)

Parameters:

array – Array-like data (anything 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 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 ak.from_jax and ak.to_numpy.