Are custom datatypes supported for Jax ?

[rishikanthc]

I found this: #5356

I would like to create a wrapper around numpy or jax arrays to support some custom operations and have it fall back to jax.numpy operations wherever it can. Is something like this possible ?

[jakevdp]

No, unfortunately there’s not currently much support for that kind of extensibility, but it’s something we’re thinking about. Can you say more about what you have in mind? e.g. would a mechanism like __jax_array__ be sufficient, or would you want JAX apis to preserve the type of the objects on output?

[rishikanthc]

For now I'm using __jax_array__ but I would like a mechanism where JAX apis could preserve the type of the objects. I'm also not clear on how jit, vmap and grad plays out with this mechanism. I have registered my custom type as a PyTree as well to make sure that jit can be used but not sure if that's sufficient.

[jakevdp]

Thanks. For what it's worth, __jax_array__ is not documented or tested, and we've never meant it to be a supported feature. Use it at your own risk (cc/ @mattjj)

[rishikanthc]

Yeah i understand the risks. I'll be on the lookout and update things when a mechanism for this rolls out. Unfortunately for my current use case __jax_array__ seems to be the way to go for now. Thanks for warning though. Appreciate the help.

[Jacob-Stevens-Haas]

@jakevdp I share OP's question, and since you asked, here's my use case. I have a numpy array subclass in a library that I'm hoping to migrate to Jax. The subclass tracks axis labels, e.g.

data = np.random.normal(size=(10, 10, 100, 3))
arr = AxesArray(data, {"ax_spatial": [0, 1], "ax_time": 2, "ax_coord": 3}

This makes some code more readable:

differentiate(arr, axis=arr.ax_time)
# alternate differentiate(arr, axis=-2)

It also makes it much easier to debug math code, since I've overloaded slicing, transformations, and linear algebraic functions to correctly handle axis labels. The comment in #5356 says that __jax_array__ does not preserve user types. Ideally there would be something like numpy's __array_function__ and __array_ufunc__ for jax arrays.

[davisyoshida]

Is __array_function__ enough for that? It seems like it wouldn't be able to propagate the names for things like transposes and reductions.

Edit: Ah nevermind, the first reference I read about it didn't actually describe what it does.

[davisyoshida]

Could you say a little bit more about what your exact use case is? My Qax transform can get you some custom behavior, although part of the point is to make it look like your object is just an ordinary array. This is basically a helper which makes writing a pretty common type of JAX transform easier.

[rishikanthc]

So I'm building a package for a neuro-symbolic computing technique. My goal is to have a custom datatype that stores information as numpy (jax / numpy) arrays. I want to overload a few operators like __add__, __sub__, __mul__ and provide a couple custom operations on them. The datatype needs to support jax.grad operations on them and allow me to define custom grad operations for the operations that I override and for any custom operations. I would like for the entire setup to be compatible with jax.jit and jax.vmap. In addition I would like for it to gracefully fallback to jax.numpy operations like concatenate, stack etc etc. So basically I want to add a wrapper around jax.numpy.ndarray that supports a few custom operations while falling back to the default functions that jax supports.

I will take a look at Qax. If this is feasible with your package, I would appreciate any pointers or documentation for doing this. I appreciate your help. Thanks a tonne ! :)

[davisyoshida]

Here's an example of a class which wraps a single array and intercepts all JAX operations. I made it so that by default it just executes the ordinary operation and then propagates the custom type. Inside any function decorated with qax.use_implicit_args, JAX operations will just see MyType instances as being an ordinary array, but you can customize how any primitive (like add, mul, transpose, etc.) is handled. The decorated function is also compatible with grad, vmap etc.

The idea of Qax is to hide away a bunch of the details of custom JAX transforms/interpreters for some common use cases. If you want even more customized behavior you could take a look at their tutorials: https://jax.readthedocs.io/en/latest/autodidax.html

from dataclasses import dataclass
import jax
import jax.numpy as jnp

import qax

@dataclass
class MyType(qax.ImplicitArray, warn_on_materialize=False):
    # You can define as many fields as you want, but in this case we're just wrapping a vanilla array
    data: jax.Array

    def materialize(self):
        # This method defines how to materialize the implicit array into a regular array if necessary
        return self.data

    @classmethod
    def default_handler(cls, primitive, *args, params=None):
        # I haven't made the "handle all primitives uniformly" case very easy to use yet,
        # but you can override the default_handler() method to customize how your type
        # handles them. See materialize_handler() in qax/implicit/implicit_array.py for a reference
        subfuns, bind_params = primitive.get_bind_params(params)

        unpacked_args = (arg.data if isinstance(arg, cls) else arg for arg in args)
        result = qax.primitives.default_handler(primitive, *unpacked_args, **bind_params)
        return MyType(data=result)

    def __str__(self):
        return f'MyType({self.data})'

# For handling specific primitives, use the primitive_handler decorator, and either pass
# an explicit primitive like jax.lax.add_p, or the name of the primitive
# The type annotations are used for dynamic dispatch with beartype, so you can also
# define additional overloads for different combinations or permutations of types.
# For commutative ops like add, the argument order will be swapped to put
# your type first, but you can disable that behavior if you want.
@qax.primitive_handler('add')
def my_custom_add(op, x : MyType, y : MyType | jax.Array):
    if isinstance(y, MyType):
        y = y.data

    return MyType(x.data + 0.5 * y)

def main():
    x = MyType(jnp.arange(5, dtype=jnp.float32))

    @qax.use_implicit_args
    def some_function(x):
        return jnp.sin(x) * 2

    print(some_function(x))
    # prints: MyType([ 0.         1.682942   1.8185949  0.28224   -1.513605 ])

    @qax.use_implicit_args
    def test_add(x):
        return (x ** 2) + 4

    print(test_add(x))
    # prints: MyType([ 2.  3.  6. 11. 18.])

    def loss_fn(x):
        # Because this class wasn't decorated with qax, it can "see"
        # that x is a MyType instance, so we can actually refer to
        # fields like x.data
        return jnp.sum(test_add(x).data)

    # MyType is just a pytree, so you can use it in transformations like grad and vmap
    print(jax.grad(loss_fn)(x))
    # prints: MyType([0. 2. 4. 6. 8.])

if __name__ == '__main__':
    main()

Currently I don't add operator overloading by default, so things like MyType(x) + 3 outside of the context of a use_implicit_args will error. You could just define an __add__ method, though, then just use + in the primitive handler.

[patrick-kidger]

+1 for Davis' recommendation of Qax as the appropriate solution to this problem.

And of course, I'll also advertise my own take on this -- Quax -- which as the name suggests is inspired by Davis' work. We've just done a 0.0.3 release and now have some shiny new documentation.

(@Jacob-Stevens-Haas -- you mention wanting named axes. As a toy implementation of exactly that, see quax.examples.named.)

[Jacob-Stevens-Haas]

Very cool, and I'll have to test it out! Enabling Jax in our package is a bit of a long term goal, and I guess I need to get more familiar with the ecosystem of ___ax packages (e.g. diffrax, quax). Not sure if quax.named allows for automatically creating named axes with slicing, e.g. arr[0, [[1, 2], [0, 1]]], but if the ecosystem is as compatible with subclasses as it seems, I'll be able to re-use my most of my custom __getitem__, __reshape__, etc with a new super.

[patrick-kidger]

The particular quax.examples.named example that's included doesn't, but that's just because it's only an example, not a full library. It should be possible to support exactly what you're after here though.

As a technical note this would be implemented by registering a rule for jax.lax.slice_p, rather than through __getitem__, btw. (This ensures compatibility with representing computations as both Python functions and JAX jaxprs.) I imagine the implementation inside would probably be pretty similar though.

[Jacob-Stevens-Haas]

Makes sense that a jax.laxative would unblock my development pipeline. Thanks Dr. Kidger and I'll reach out if there's any side effects.

... sorry in advance if you get a lot of constipation jokes for working on jax.lax