Show HN: Equinox (1.3k stars), a JAX library for neural networks and sciML

1 point

3 years ago

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` namespace, libraries like Equinox are then akin to the `torch.nn` namespace.

Because of its speed and features, right now JAX+Equinox is my favourite approach to numerical computing. So I'd love for some more people to try it. What do you think?