coordax.NDArray

class coordax.NDArray

Abstract base class for non-JAX arrays that can be used with Coordax.

THIS FEATURE IS EXPERIMENTAL, AND THE DETAILS ARE SUBJECT TO CHANGE. We currently use it only to support jax-datetime, but expect it might be useful in the future for other array types, e.g., with arbitrary units.

To register a new NDArray, use coordax.experimental.register_ndarray().

NDArray instances are also expected to be registered as JAX pytree nodes: https://docs.jax.dev/en/latest/pytrees.html#extending-pytrees

__init__()

Methods

__init__()
transpose(axes)	Transpose this array to this given axis order.

Attributes

ndim	Number of dimensions in this array, typcially len(self.shape).
shape	Shape of this array.
size	Size of this array, typically math.prod(self.shape).