ak.from_jax

Defined in awkward.operations.convert on line 545.

ak.from_jax(array, regulararray=False, highlevel=True, behavior=None)
Parameters
  • array (jax.numpy.array) – The jax.numpy.array to convert into an Awkward Array.

  • regulararray (bool) – If True and the array is multidimensional, the dimensions are represented by nested ak.layout.RegularArray nodes; if False and the array is multidimensional, the dimensions are represented by a multivalued ak.layout.NumpyArray.shape. If the array is one-dimensional, this has no effect.

  • highlevel (bool) – If True, return an ak.Array; otherwise, return a low-level ak.layout.Content subclass.

  • behavior (None or dict) – Custom ak.behavior for the output array, if high-level.

Converts a JAX array into an Awkward Array.

The resulting layout may involve the following ak.layout.Content types (only):

  • ak.layout.NumpyArray

  • ak.layout.RegularArray if regulararray=True.

See also ak.from_cupy and ak.from_numpy.