search

lab-cosmo / sparse_accumulation

11 stars 1 forks source link

jax support #11

Open mariogeiger opened 1 year ago

mariogeiger commented 1 year ago

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 commented 1 year ago

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.