equinox

Why we included this project

Equinox fills the gap in JAX for people who want to define neural networks without fighting the framework. Models are just ordinary Python classes registered as PyTrees, so the syntax reads like PyTorch while everything stays fully interoperable with plain JAX, including jit, grad, and vmap. That design removes a lot of the boilerplate around stateful parameters and custom tree handling. It also includes filtered transformation APIs and runtime error checking, which are handy for research code where you need to keep things debuggable and hackable. If you do deep learning or differentiable scientific computing in JAX and want a lighter alternative to Flax or Haiku, this is worth a look.