Can we inherit from jax.Tensor ?

[samsja]

Context

What is the best way to inherit from jax.Tensor ?

I want to create a wrapper around ´jax.Tensor` that add on our two methods to it.

I mainly want to things:

• That my new class is still a jax.Tensor or at least that I can still be able to use it with all of the jax ecosystem
• That I can transform any jax.Tensor to my class without any memory duplication.

With Pytorch and numpy I can do that easily by just subclassing the tensor types but with Jax it is not working. My guess is that it is because the real class of a jax.Tensor change with time depending on the operation that we do on it.

So my question is, what is the best way to achieve these two points in Jax ?

[samsja]

cc @agaraman0

[davisyoshida]

Can you give an example use case? I'm working on a convenience class which is adjacent to this, so it might be that you can get what you want although not in exactly the same way.

[agaraman0]

instead of jax.Tensor it's jax.Array and we are trying to inherit it by doing something like

class WrapperClass(jax.Array):
    pass

but it does not work as expected as jax.Array is an interface class and jax.ArrayImpl can not be inherited because it's base class and subclassing is not allowed on the base class. so we are looking for a way to inherit Jax's native data type class by which its native methods and properties can be inherited in WrapperClass

cc: @samsja

[jakevdp]

No, there is no way to subclass JAX arrays in such a way that they will be compatible with JAX transformations.

This may be an XY problem – if you can tell us more about your actual end-goal, we may be able to recommend a good approach that is possible within JAX's API & computational model.

[samsja]

FYI: we are maintainer of docarray and we are trying to add support for jax, currently we support pytorch, numpy and tf.

What we want is to be able to add a couple of function to our tensor class (just some helper function). Thus we would like to be able to extend the Jax.Array object, tho it is understandable why it is not working but we were wondering if they were solutions.

Here is an example of how it looks like for pytorch tensor : https://github.com/docarray/docarray/blob/main/docarray/typing/tensor/torch_tensor.py

[jakevdp]

I see, thanks. There is not any easy way to create a JAX-compatible array with additional methods defined.

I wonder if you could do something different, e.g. define functions in docarray that accept JAX arrays as inputs, rather than defining new methods on the array objects?

[samsja]

We have a possible workaround already, doing composition over inheritance, but this is a bit cumbersome for users because it is a bit more verbose to access the pure tensor directly.

Alternatively is there smth like __torch_function__ and __array_function__ in jax that could make jax function work with non Jax class easier ?

[jakevdp]

No, JAX doesn't provide any public mechanism for inheriting from the array class or extending its API. This has not been a design goal of the project in the past. Creating a way to do so that is compatible with all JAX transformations would be a non-trivial API design and a fairly major project.

Typically, the alternative approaches have involved defining transformations that allow user-specified objects to work with JAX programs (similar to the jax.experimental.sparse.sparsify transformation, for example). That's possible, but I suspect it's somewhat of a heavier solution than what you have in mind.

Do you think this is something you can work around?

[samsja]

yes we can work around no worries.

I would say as a user it feels more natural to be able to extend torch tensor and numpy array. But I totally understand that this is not in the scope of JAX for now and that they are reason to not allow it.

Thanks for the quick help 😃