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.

[martinjrobins]

@jackbmontgomery any updates on this?

[jackbmontgomery]

Hi @martinjrobins, sorry for the lack of communication. I started my masters so I wasn't able to work on this very much. It became clear that I would need C++ and I do not know the language. I replicated the jax.ffi example with rust here but I did not have time to do anything more than that.

I see that jaxbind updated their backend to use the new ffi api. So as soon as I get some free time I will try adapt that to make the jax wrapper but I don't think I will have enough time for this until June next year when I finish.

[martinjrobins]

Copying some info on the javascript wrapper from #203:

Sure! diffsol will work with wasm out of the box in the way you are using it, but I'm also keen to get a proper javascript wrapper library, similar to pydiffsol in Python (https://github.com/alexallmont/pydiffsol). The main TODOs here are:

1. Write an initial javascript wrapper using https://github.com/wasm-bindgen/wasm-bindgen. This can use callbacks into javascript code so that the user can provide rhs, mass matrix functions as plain javascript functions. This is slow, but does not require you to build diffsl and add it to your bundle.
2. Next step is to add optional diffsl support. This requires you to onstruct a suitable compilation toolchain to be able to build llvm, enzyme, diffsl and diffsol using wasm. This requires building llvm and enzyme to wasm and linking them into a cargo build, I've made a start on the llvm and enzyme build here: https://github.com/martinjrobins/wasm-compilers. This requires llvm v20, and I'm in the process of adding support for this here (martinjrobins/diffsl@4ffa31f).

[jackbmontgomery]

Hi @martinjrobins,

I have got an update on the JAX wrapper I am writing here. I have the jit and grad working and I have included some benchmarks in the repo as well.

The biggest architectural decision is a translation between jaxpr, the intermediate representation for jax, and diffsl. It just allows a user to define a normal python function (with some constraints) for the rhs of the ode, generate the diffsl, then sent that to diffsol though a small C++ shim.

I will continue to work on this. Specifically, adding other solvers which should be easy, and adding to the translation between jaxpr and diffsl.

Please let me know if you have any thoughts or comments.

PS. I also have not developed any open-source software before but I am wanting to start so any advice on project structure, documentation, etc. would be greatly appreciated

[martinjrobins]

Thanks @jackbmontgomery, this is looking great! Good work in getting all this to work and lowering the python to diffsl. What are the constraits? does it work with jax arrays?

In terms of adding more solves, matrix, linear solver and scalar types, your best bet is to use the diffsol-c crate, this uses a "poor mans dynamic dispatch" (i.e. lots of match statements!) to allow you to select all those at runtime using enums. Have a look at this PR in pydiffsol to see how it works: alexallmont/pydiffsol#78.

I see your point about the lifetime of the checkpointing struct and the problem, that is an issue and I'll look into the best way to handle this, its only if you have a huge problem and are worried about running out of memory that you actually need that lifetime at all.

I would encourage you to make a start on the docs (probably using sphinx), ci and in particular getting the wheels built in ci in different platforms (linux/mac/windows). Again you can use what Alex has setup in pydiffsol as a guide on how to do this.

[jackbmontgomery]

Hi @martinjrobins,

Thanks! On constraints: there's no control flow (no loops or conditionals), and a few matrix operations I haven't lowered yet eg. reductions, concatenation, dot_general. Nothing I'd call a blocker so far.

Yes, it works with JAX arrays throughout. The parameters and initial conditions are JAX arrays, and the RHS function can use the usual elementwise ops and they all lower correctly.

On diffsol-c: I can't use it at the moment because I rely on jax.custom_vjp, which needs to propagate arbitrary cotangents through the adjoint solve. The current public adjoint API (solve_sum_squares_adj) hardcodes a sum-of-squares loss, so it doesn't fit. If diffsol-c exposed a general adjoint something like the existing pub(crate) solve_adjoint_backwards that takes dgdu directly then I could migrate onto it. Happy to contribute that if it'd be useful.

Docs, CI, and cross-platform wheels all make sense as next steps. I'll look at Alex's pydiffsol setup.

[martinjrobins]

If diffsol-c exposed a general adjoint something like the existing pub(crate) solve_adjoint_backwards that takes dgdu directly then I could migrate onto it. Happy to contribute that if it'd be useful.

Yea agreed, that would be a nice addition, the main blocker on this I suspect would be the lifetime on the checkpointing struct, I'll look at what I can do to remove/get around this. If you have any ideas please let me know, very happy to take contributions!