Wanted - Developers for higher-level language wrappers

[martinjrobins]

Diffsol is designed to be easy to use from higher-level languages like Python or R. I'd prefer not to split my focus away from the core library, so I'm looking for developers who would like to lead the development of these wrappers. If you're interested, please get in touch.

• Python (e.g. using PyO3). https://github.com/alexallmont/pydiffsol.
• Python ML frameworks (e.g. JAX, PyTorch)
• R (e.g. using extendr).
• Julia
• Matlab
• JavaScript backend (e.g using Neon)
• JavaScript in browser (e.g. using wasm-pack)
• Others, feel free to suggest your favourite language.

[CGMossa]

Hello Martin! I would love to write the R-package for diffsol. I have been lurking (and part of the discord) just for this moment! Please feel free to communicate with me about your plans, etc.

[jackbmontgomery]

Hi Martin, I am interested in developing a high level library for the python ml frameworks!

[martinjrobins]

Hi Martin, I am interested in developing a high level library for the python ml frameworks!

Great, thanks @jackbmontgomery. FYI I was thinking that this would probably involve wrapping diffsol as a custom JAX operator. You can read more about jax's FFI interface here (https://docs.jax.dev/en/latest/ffi.html), and I notice that there are already rust bindings for XLA here (https://docs.rs/xla/latest/xla/) which might be useful. You will be able to use the forward and adjoint sensitivity functionality in diffsol to implement the forwards and backwards gradient passes required by JAX. I would implement this as a separate python library using pyo3 and maturin

This is just my idea, but feel free to propose anything else that you are interested in. Please feel free to lead this project however you want, and I'll be available for help and advice as needed. Once you have something I can link to it from this repo and the documentation as well. I'm happy to jump on a call as well if that is useful for you.

[jackbmontgomery]

That all makes sense to me @martinjrobins. I have not used jax's FFI feature before so let me start there. Thanks a lot for the suggestions and I may be in contact soon with some questions.

[jackbmontgomery]

Hi @martinjrobins, I am just wanting to follow up on what I have been reading and working on with the JAX wrapper. The rust XLA bindings do not include the macros or the buffers that are required by XLA for the arguments and results in custom calls (https://openxla.org/xla/custom_call). So I don't think the crate will work out of the box for this. I have thought of some options for what to do next:

1. Build on the xla crate to implement these features that are needed by XLA for foreign calls. And then proceed as you have described with py03 and maturin.

2. Use C++ as the interface between Rust and JAX with c++ binding created by cxx (https://cxx.rs/). Then using nanobind to generate the python package and stubs. This is what I have been experimenting with to replicate the tutorial in this repo.

3. The final option I have seen is to use a python package called jaxbind. This package allows you to create jax primitive operations by defining the jvp, vjp and behaviour for jax's abstract evaluation to create the operator. This will allow for the use of py03 and maturin as well.

I’d love to hear your thoughts on any of these approaches - whether you see any issues, or if you have preferences or suggestions. Any feedback would be greatly appreciated!

[martinjrobins]

you would have to experiment, but I suspect 2. is the most straightforward and most likely to succeed and the one I would recommend you start with (as you already have), while 1 would be more work but would provide most benifit to the wider rust community. 3 is a bit risky because I don't know the details of jaxbind, and what its restrictions and drawbacks are (e.g. do you have to go back to python with every call to f or vjp/jvp?)

[jackbmontgomery]

Yes I was thinking the same. I think what I will do is carry on working along 2 and experiment with 1 since I think that would be the best implementation of the jax wrapper.