equinox

Anecdotally, I used diffrax (and equinox) throughout last year after jumping between a few differential equation solvers in Python, for a project based on Dynamic Field Theory [1]. I only scratched the surface, but so far, it's been a pleasure to use, and it's quite fast. It also introduced me to equinox [2], by the same author, which I'm using to get the JAX-friendly equivalent of dataclasses. `vmap`-able differential equation solving is really cool. [1]: https://dynamicfieldtheory.org/ [2]: https://github.com/patrick-kidger/equinox

There's quite a few other libraries associated with Equinox in the JAX ecosystem: https://github.com/patrick-kidger/equinox?tab=readme-ov-file... I've enjoyed using Equinox and Diffrax for performing ODE simulations. To my knowledge the only other peer library with similar capabilities is the Julia DifferentialEquations.jl package.

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

If you like PyTorch then you might like Equinox, by the way. (https://github.com/patrick-kidger/equinox ; 1.4k GitHub stars now!) Basically designed to offer PyTorch-like syntax for working with JAX. The latter is excellent for the reasons the sibling replies have stated, but PyTorch absolutely got the usability story correct.

You can do JAX->TF->ONNX I believe: https://github.com/patrick-kidger/equinox/blob/main/equinox/...

I'm going to disagree here! Classes and functional programming can go very well together, just don't expect to do in-place mutation. (I.e. OO-style programming.) You might like Equinox (https://github.com/patrick-kidger/equinox ; 1.4k GitHub stars) which deliberately offers a very PyTorch-like feel for JAX. Regarding speed, I would strongly recommend JAX over PyTorch for SciComp. The XLA compiler seems to be much more effective for such use cases.

I feel like I see the opposite -- that everything scientific computing is getting rewritten in something autodifferentiable! Whether that's JAX or something else. My experience might be biased though: shameless advert for Equinox (https://github.com/patrick-kidger/equinox, 1.4k GitHub stars), which is now the foundation of quite a lot of SciComp in JAX. (Both internal and open-source.)

Hey folks! I wanted to advertise Equinox -- my now-surprisingly-popular ( :D ) JAX library for numerical models. These days that often means "neural networks", but I like to emphasise that this also includes ODEs/SDEs/linear solves, etc. For those already using JAX, then Equinox is interesting because (a) it ships with a NN library, and (b) this is built around the idea that "everything is a pytree", which makes things easy to reason about and easy to compose. Furthermore (c) Equinox offers advanced tools like true runtime errors, out-of-place pytree surgery, and checkpointed while loops, and AFAIK in the JAX ecosystem these are unique to Equinox. For those most familiar with PyTorch: for many use cases (sciML in particular), JAX has a much stronger compiler, more advanced autodiff, etc. And whilst JAX itself is akin to the `torch` namesp

You're thinking of `jax.closure_convert`. :) (Although technically that works by tracing and extracting all constants from the jaxpr, rather than introspecting the function's closure cells -- it sounds like your trick is the latter.) When you discuss dynamic allocation, I'm guessing you're mainly referring to not being able to backprop through `jax.lax.while_loop`. If so, you might find `equinox.internal.while_loop` interesting, which is an unbounded while loop that you can backprop through! The secret sauce is to use a treeverse-style checkpointing scheme. https://github.com/patrick-kidger/equinox/blob/f95a8ba13fb35...

Shameless advert -- Equinox is a neural network library for JAX based entirely around pytrees: https://github.com/patrick-kidger/equinox (Now on 1.1k stars so it's achieved some popularity!) This makes model-building elegant (IMO), without any new abstractions to learn. Quite a PyTorch-like experience overall.

It sounds like you're concerned about how downstream libraries tend to wrap JAX transformations to handle their own thing? (E.g. `haiku.grad`.) If so, then allow me to make my usual advert here for Equinox: https://github.com/patrick-kidger/equinox This actually works with JAX's native transformations. (There's no `equinox.vmap` for example.) On higher-order functions more generally, Equinox offers a way to control these quite carefully, by making ubiquitous use of callables that are also pytrees. E.g. a neural network is both a callable in that it has a forward pass, and a pytree in that it records its parameters in its tree structure.

+1 for JAX. Basically designed to be the successor to TensorFlow, and much nicer to work with. Strangely I've not seen it discussed around HN much but it's what I do 100% of my work in these days. Whilst I'm here: shameless self-promotion for Equinox and Diffrax: https://github.com/patrick-kidger/equinox https://github.com/patrick-kidger/diffrax Which are neural network and differential equation libraries for JAX. [Obligatory I-am-a-googler-my-opinions-do-not-represent-your-employer...]

Autodifferentiable programming! Neural networks are the famous example of this, of course -- but this can be extended to all of scientific computing. ODE/SDE solvers, root-finding algorithms, LQP, molecular dynamics, ... These days I'm doing all my work in JAX. (E.g. see Equinox or Diffrax: https://github.com/patrick-kidger/equinox, https://github.com/patrick-kidger/diffrax). A lot of modern work is now based around hybridising such techniques with neural networks. I'd really encourage anyone interested to learn how JAX works under-the-hood as well. (Look up "autodidax") Lots of clever/novel ideas in its design.