Level 1 Extensibility: polymorphic inputs#

What I’ll call “Level 1” extensibility is the desire that JAX APIs accept polymorphic inputs.That is, a user desires behavior like this:

classCustomArray:data:numpy.ndarray...x=CustomArray(np.arange(5))result=jnp.sin(x)# Converts `x` to JAX array and returns a JAX array

Under this extensibility model, JAX functions would accept CustomArray objects as inputs,implicitly converting them tojax.Array objects for the sake of computation.This is similar to the functionality offered by NumPy via the__array__ method, and inJAX (in many but not all cases) via the__jax_array__ method.

This is the mode of extensibility that has been requested by the maintainers offlax.nnxand others. The current implementation is also used by JAX internally for the case ofsymbolic dimensions.

Level 2 extensibility: polymorphic outputs#

What I’ll call “Level 2” extensibility is the desire that JAX APIs should not only acceptpolymorphic inputs, but also wrap outputs to match the class of the input.That is, a user desires behavior like this:

classCustomArray:data:numpy.ndarray...x=CustomArray(np.arange(5))result=jnp.sin(x)# returns a new CustomArray

Under this extensibility model, JAX functions would not only accept custom objectsas inputs, but have some protocol to determine how to correctly re-wrap outputs withthe same class. In NumPy, this sort of functionality is offered in varying degrees bythe special__array_ufunc__,__array_wrap__, and__array_function__ protocols,which allow user-defined objects to customize how NumPy API functions operate onarbitrary inputs and map input types to outputs.JAX does not currently have any equivalent to these interfaces in NumPy.

This is the mode of extensibility that has been requested by the maintainers ofkeras,among others.

Level 3 extensibility: subclassingArray#

What I’ll call “Level 3” extensibility is the desire that the JAX array object itselfcould be subclassable. NumPy provides some APIs that allow this(seeSubclassing ndarray) butthis sort of approach would take some extra thought in JAX due to the need forrepresenting array objects abstractly via tracing.

This mode of extensibility has occasionally been requested by users who want to addspecial metadata to JAX arrays, such as units of measurement.

Priority of__jax_array__ vs. PyTree flattening#

JAX already has a supported mechanism for registering custom objects, namely pytreeregistration (seeCustom pytree nodes).If we also supportjax_array, which one should take precedence?

To put this more concretely, what should be the result of this code?

@jax.jitdeff(x):print("is JAX array:",isinstance(x,jax.Array))f(CustomArray(...))

If we choose to prioritize__jax_array__ at the JIT boundary, then the output of thisfunction would be:

isJAXarray:True

That is, at the JIT boundary, theCustomArray object would be converted into a__jax_array__, and its shape and dtype would be used to construct a standard JAXtracer for the function.

If we choose to prioritize pytree flattening at the JIT boundary, then the output ofthis function would be:

type(x)=CustomArray

That is, at the JIT boundary, theCustomArray object is flattened, and then unflattenedbefore being passed to the JIT-compiled function for tracing. IfCustomArray has beenregistered as a pytree, it will generally contain traced arrays as its attributes, andwhen x is passed to any JAX API that supports__jax_array__, these traced attributeswill be converted to a single traced array according to the logic specified in the method.

There are deeper consequences here for how other transformations like vmap and grad workwhen encountering custom objects: for example, if we prioritize pytree flattening, vmapwould operate over the dimensions of the flattened contents of the custom object, whileif we prioritize__jax_array__, vmap would operate over the converted array dimensions.

This also has consequences when it comes to JIT invariance: consider a function like this:

deff(x):ifisinstance(x,CustomArray):returnx.custom_method()else:# do something else...result1=f(x)result2=jax.jit(f)(x)

Ifjit consumesx via pytree flattening, the results should agree for a well-specifiedflattening rule. Ifjit consumesx via__jax_array__, the results will differ becausex is no longer a CustomArray within the JIT-compiled version of the function.

Which APIs should support__jax_array__?#

JAX has a number of different levels of API, from the level of explicit primitive binding(e.g.jax.lax.add_p.bind(x,y)) to thejax.lax APIs (e.g.jax.lax.add(x,y)) to thejax.numpy APIs (e.g.jax.numpy.add(x,y)). Which of these API categories should handleimplicit conversion via__jax_array__?

In order to limit the scope of the change and the required testing, I propose that__jax_array__only be explicitly supported injax.numpy APIs: after all, it is inspired by the__array__protocol which is supported by the NumPy package. We could always expand this in the future tojax.lax APIs if needed.

This is in line with the current state of the package, where__jax_array__ handling is mainlywithin the input validation utilities used byjax.numpy APIs.