diffrax

I wrote a JAX-based neural network library (Equinox [1]) and numerical differential equation solving library (Diffrax [2]). At the time I was just exploring some new research ideas in numerics -- and frankly, procrastinating from writing up my PhD thesis! But then one of the teams at Google starting using them, so they offered me a job to keep developing them for their needs. Plus I'd get to work in biotech, which was a big interest of mine. This was a clear dream job offer, so I accepted. Since then both have grown steadily in popularity (~2.6k GitHub stars) and now see pretty widespread use! I've since started writing several other JAX libraries and we've turned this into a bit of a foundation for a JAX sciML ecosystem. [1] https://github.com/patrick-kidger/equinox [2] https://github.com/patrick-kidger/diffrax

Hi everyone! I wanted to advertise my new JAX optimisation library Optimistix! Optimistix has high-level APIs for minimisation, least-squares, root-finding, and fixed-point iteration and was written to take care of these kinds of subroutines in Diffrax (https://github.com/patrick-kidger/diffrax). Here is the GitHub: https://github.com/patrick-kidger/optimistix The elevator pitch is Optimistix is really fast, especially to compile. It plays nicely with Optax for first-order gradient-based methods, and takes a lot of design inspiration from Equinox (https://github.com/patrick-kidger/equinox), representing the state of all the solvers as standard JAX PyTrees. For those familiar with classical nonlinear unconstrained optimisation, Optimistix does some pretty nifty new things. It introduces new abstractions for modular optimisers, allowing us