Is there a way to subclass jnp.ndarray/jax.Array to add a custom field

[hamzamerzic]

We have a general purpose codebase which we would like to extend in a way that allows tagging our parameters (jnp.arrays) with metadata. A minimally intrusive solution both for implementation and typing purposes would be a solution that simply extends jnp.array object with an additional field. This is supported in NumPy and they provide detailed documentation on how to do so [link]. I wonder if there is a similar possibility in JAX?

Alternatively, is there a way to create an object that for all intents and purposes in JAX behaves exactly like a jnp.ndarray? That would solve the implementation part of things and I'm sure that we could do some typing magic to make this object look like a jax.Array.

[patrick-kidger]

This isn't possible in JAX. This is very much by design: JAX arrays need to get swapped out with Tracers and other non-array-objects internally, when JIT-ing/differentiating/vmap'ing etc.

Fortunately it's still pretty easy to carry additional metadata around, just by wrapping your array into a pytree. For this example I'll use Equinox for easy handling of pytrees, but if you prefer it's possible to do something analogous by registering your own custom pytrees.

import equinox as eqx
import jax
import jax.numpy as jnp
import jax.tree_util as jtu
from typing import Any

class HasMetadata(eqx.Module):
    array: jnp.ndarray
    metadata: Any

def is_metadata(x):
    return isinstance(x, HasMetadata)

def resolve_metadata(x):
    if is_metadata(x):
        # Or if you wish, do some other processing based on the value of the metadata.
        return x.array
    else:
        return x

linear = eqx.nn.Linear(...)
meta_weight = HasMetdata(linear.weight, {"your": "metadata"})
linear2 = eqx.tree_at(lambda x: x.weight, linear, meta_weight)
# linear2 has metadata annotated on its weight

@jax.jit
@jax.grad
def loss(model, x, y):
    model = jtu.tree_map(resolve_metadata, model, is_leaf=is_metdata)
    pred_y = model(x)
    return jnp.mean((y - pred_y)**2)

loss(linear2, ...)

You can also use the same trick for e.g. creating linear layers with symmetric weight matrices: make the "resolution" step return 0.5 * (x.array + x.array.T) instead of just x.array.

[hamzamerzic]

Thanks for the answer Patrick! We are aware of this approach and is indeed what we used in our initial implementation. The problem is that we have a lot of code that does things like x.astype(jnp.float32), which we now need to replace with a tree_map which is not really feasible for a general purpose library.

We've thought about this a lot and the only solution we see is to be able to subclass an array to add some extra fields.

[jakevdp]

Hi - sorry this is ending up being a challenge. I'll echo Patrick's statement: JAX doesn't support straight subclassing of arrays, and is unlikely to ever do so.

One point about the "just add attributes" approach: it leads to a lot of design challenges, e.g.

x.special_attribute = 'A'
y.special_attribute = 'B'
z = x + y
assert z.special_attribute == ???

It's not clear what the ??? should be in this case, and binary ufuncs are one of the fundamental units of building a computation! I could imagine ??? being many different things depending on the application. It's one reason that adding custom attributes isn't really feasible for a general-purpose library like JAX, even if it were technically possible (which, given JAX's current design choices, it is probably not). Does that make sense?

One alternative that's a bit more heavy-weight: you could define your pytree-based array-like object with metadata, and define a custom transform that has implementations of every JAX primitive that specifies what to do in situations like this (binary operations between objects with different attributes). Then you wrap all your code in this attribute transform and write code as normal. If that's something you want to pursue, I could try to point you in the right direction.