Extending JAX with Rust

[jeertmans]

Hi everyone!

I really enjoy JAX, and I also really appreciate programming in Rust, and I would love to be able to use Rust code to implement some performance-critical functions and export them to Python.

Currently, I do it using PyO3 and rust-numpy, but I lose all the benefits from XLA (mainly automatic differentiation for me), so I eventually have to first convert NumPy arrays to JAX arrays in Python.

I am aware of Extending JAX with custom C++ and CUDA code and Custom operations for GPUs with C++ and CUDA, which are two great tutorials, but I feel there should be something easier when it comes to writing a JAX extension with Rust, especially when tools and libraries like PyO3, maturin, rust-numpy, and xla-rs exist.

Have some of you already tried implementing a JAX (or XLA) extension with Rust?

I was considering giving it a try myself, but I am of course interested to see if other people have tried before me.

[dfm]

I did play around with this a bit at one point, but I'm not aware of any worked examples. I'd recommend taking a look at the new foreign function interface. I expect the easiest approach would be to write a C++ shim that uses the FFI to dispatch to Rust, rather than trying to reimplement the FFI in Rust directly. Using variadic arguments and results you could probably write a pretty general purpose interface. Hope this helps!

[jeertmans]

Thanks! I did already check about the FFI, but I was worried that this would just require to re-implement code that was already sort-of present in xla-rs.

[jeertmans]

I did play a bit with the FFI this weekend. My plan was to generate FFI bindings (using autocxx and autocxx-build) from the headers in jax.extend.include_dir(), but some important C++ types seems to be impossible to translate in Rust (e.g., Buffer, or Binding), at least using cxx, and also macros cannot be exported, which means macros like XLA_FFI_DEFINE_HANDLER_SYMBOL must be rewritten manually in Rust, which defeats a bit the purpose of using automated tools like autocxx.

As you suggested, I guess the only realistic solution is to actually call Rust code from C++, but my hope is still to manage to use PyO3 (or at least Maturin) to automate the packaging, because it is (in my opinion) a major advantage of Rust over C/C++ when it comes to creating native extensions for Python.

I will let you know if I manage to have something working.

[dfm]

Thanks for the update! Yeah, I expected that that might end up being an issue. Please do keep this thread updated as you make more progress!

[jeertmans]

Hi @dfm! For your information, I made some "advancements" today, as I have some sort of prototype code that compiles, but it triggers a segmentation faults relatively hard to trace (didn't dig very deep tho).

So it's not working at the moment, but I think the project structure is close to something that should work, if you want to take a look: https://github.com/jeertmans/extending-jax.

Otherwise, I will hopefully find a fix, next time I have some time to work on this :-)

[aseyboldt]

This is something I'm also very interested in. Mostly the other way round however. I would like to be able to call a jax compiled function from rust. The way I do it right now is to just call a python function which calls jax, but this adds a lot of overhead (including running into the GIL).
I could use https://github.com/LaurentMazare/xla-rs, but with this I have to somehow distribute xla again, even though it is already used (but I think not really exposed as a shared library) in jaxlib.

[jeertmans]

I will probably have some time to try going from Rust to Python starting November, and will update here if I have any news