Linear Operator implementation and einsum

[quattro]

Hi all, I'm working on an implementation of a simple LinearOperator class in JAX that I'm using to represent a centered and scaled sparse matrix. The gist of the class is along the lines of

from jax import Array
from jax.typing import ArrayLike

class LinearOp:
    def __mat__(self, x: ArrayLike) -> Array:
       [...]
    def __rmat__(self, x: ArrayLike) -> Array:
       [...]
    def transpose(self):
       return LinearOpT(...)

    T = property(transpose)

    @propery
    def shape(self) -> Array:
        [...]

    # other properties like ndim, size, etc...

class LinearOpT(LinearOp):
    [...]

It seems to be working out well for my needs so far using basic scaling and matrix multiplication, but I'd like to extend it to be compatible with calls to the einsum function. The idea would be to leverage einsum and jax.Array matrix mult/scaling seamlessly to achieve something like,

my_data = LinearOp(...)
X, Y = # some jax arrays
X @ my_data # works
my_data @ Y # works

A, B = # more jax arrays
A @ my_data.T # works
my_data.T @ B # works

jnp.einsum("np,nk,kp->p", X, my_data, Y) # doesn't work

What is not clear to me is what additional private methods need to be defined in order to work with einsum. The documentation on extending numpy.ndarray discusses ufunc and other attributes, but it's unclear how much of those design decisions extend to the jax.Array setting. I would greatly appreciate any direction on this matter.

Cheers!

[mattjj]

Thanks for the question!

That seems like a cool thing to build. Unfortunately JAX doesn't currently offer a good way to overload jax.numpy function application (and the NumPy's overriding mechanisms probably won't work here).

The best current solution I can think of is just to define your own NumPy wrapper API in a package, like import quattro.numpy as qnp, and then write functions calling qnp (or call it jnp or np if you prefer).

We've talked about adding a NumPy-overriding API to JAX, but it's never risen to the top of the priority list.

What do you think?

[mattjj]

In your above example, where does the einsum name come from?

[quattro]

jax.numpy.einsum

[quattro]

Hi @mattjj , I've gone ahead and edited the example code to be a bit clearer on my intention using type annotations, along with a few examples of what works, versus what doesn't. The end goal would be to throw in LinearOp into jnp.einsum function calls seamlessly. I'm unsure if I need to add __jax_array__ to the LinearOp definition, or some other internal function to get it to play nicely.

[jakevdp]

As @mattjj said, there's not currently any way to do what you're asking. We've thought about supporting this kind of extensibilty for the JAX API, but no such support has been implemented.

Your best option at the moment is probably to make your own numpy-stype API that mirrors JAX, so for example you could define an einsum function that accepts LinearOp types and does the appropriate thing with them. I realize that's not a very satisfying response to your question.

A more involved possibility would be to define a custom JAX transform that transforms JAX code into something that would accept LinearOp types. There are some examples to get you started here: https://jax.readthedocs.io/en/latest/notebooks/Writing_custom_interpreters_in_Jax.html

[quattro]

I see, thanks for the explanation. I appreciate you both taking the time to respond. Cheers!