ak.to_jax¶
Defined in awkward.operations.convert on line 593.
-
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
.