jax support

[mariogeiger]

Hi,

I'm the author of e3nn-jax and I was wondering if you would be motivated to also provide a jax interface for your code such that I can make your code an optional dependency of e3nn-jax to compute the tensor product, (maybe also spherical harmonics and symmetric contractions...)

What do you think?

Best

[Luthaf]

This package contains custom CUDA code, and this is what makes it relatively fast. From my understanding, it is still quite hard to register custom CUDA kernels with JAX (although they are working on it). I agree that a JAX version would be interesting, but it looks like a lot of work right now.