JAX support

[teddykoker]

The README mentions "future frontend support outside of torch"; are there any plans for integration with JAX, e.g. using foreign function interface?

[asglover]

Hello Teddy,

It's great to hear from you! That's the idea! I have not started working towards it yet, as the work is not yet finished on the torch integration.

If / when we start work on JAX bindings, is there a project you have in mind, where you would be interested in integrating it?

[teddykoker]

I'll ping you when I start working on JAX bindings.

Thanks, appreciate this!

Could you make a list of the all the JAX features you require / use?

Currently we just need the uvu tensor product and the fused convolution. Happy to provide more information though.

Any chance you could open discussions on nequix?

Absolutely! Just started the discussion and added some reference numbers across jax/torch/torch + oeq.

[asglover]

Could you make a list of the all the JAX features you require / use?

Currently we just need the uvu tensor product and the fused convolution. Happy to provide more information though.

I meant more along the lines of, what JAX features are you using. So for torch, you are using torch.compile, autograd, .to, and not using torch.jit.script, autocast ect. Just because I want to know for planning if you use some particularly tricky jax feature, like distributed.

[teddykoker]

Ah yes -- the JAX implementation uses equinox, which is a very lightweight library around JAX. As long as it supports the usual jax.jit(), jax.grad() it should be great. We do use jax.pmap() for training on multiple gpus as well, but I'm not sure if that would add any difficulty to the kernel implementation.

[asglover]

Ah, nice, that is the exact kind of thing I was wondering about. There is no documentation about the interaction of FFI and pmap, but it's the most benign kind of parallelism, so I'm guessing that pmap will fully compose with FFI? Thanks for the info!

[teddykoker]

Yes I could not find anything either - I'm also assuming that because it is simple data parallelism it should simply compose with FFI, but I guess we shall see!

[abhijeetgangan]

Is there any update on the JAX integration? Would be nice to also use it with the mace-jax. I also plan to integrate these models for simulation in JAX-MD and having these kernels would be great!

[asglover]

Hi Abhijeet,
No meaningful update, there are still other things in the queue before JAX work will begin, but I will post here when it does. At that time if you are still interested in integrating the kernels, maybe we can collaborate.

• Austin

[abhijeetgangan]

Looking forward to it. Happy to collaborate!