Equinox is a Python library that simplifies building and training neural networks and performing scientific computing within JAX. It provides a PyTorch-like class-based API while maintaining compatibility with JAX's functional programming paradigm and its ecosystem. Equinox isn't a framework; instead, it offers tools for filtered transformations and PyTree manipulation, allowing for fine-grained control over models. It is currently at version 0.13.6 and is actively maintained.
pip install equinoxVerified import paths — ran on the pinned version, not inferred.
This quickstart defines a simple Multi-Layer Perceptron (MLP) using `eqx.Module` and `eqx.nn.Linear`. It demonstrates how to build a neural network with a PyTorch-like class syntax and then use it for inference. The model's parameters are initialized using JAX random keys.
Always use `eqx.filter_jit` and `eqx.filter_grad` for transformations on Equinox models, or explicitly `eqx.partition` your model into `(trainable, nontrainable)` parts before applying raw JAX transformations.
Embrace functional programming paradigms. To update parts of a model, use functional approaches like `eqx.tree_at` to create a new model with updated values rather than modifying attributes directly. For example: `model = eqx.tree_at(lambda m: m.attribute, model, new_value)`.
Be mindful of the `deterministic` argument's precedence. Explicitly pass `deterministic=True` or `deterministic=False` at call time for clarity, especially during inference or training phases, if you want to override the default behavior set at initialization.
Use Equinox's filtered transformations, such as `equinox.filter_jit` or `equinox.filter_vmap`, instead of their `jax` counterparts. These filtered versions automatically treat non-array leaves as static.
Ensure that a `jax.random.PRNGKey` is passed as the `key` argument when instantiating such modules: `linear = eqx.nn.Linear(in_size, out_size, key=jax.random.PRNGKey(0))`.
Define all attributes of an `equinox.Module` using class annotations. For initialization, use `eqx.field` for default values or metadata, or set them directly in `__init__` for values dependent on constructor arguments, without reassigning attributes that are already part of the module's PyTree.
When using `jax.lax.scan` with an `equinox.Module`, wrap the module function in a `lambda` to ensure it is treated correctly, for example: `jax.lax.scan(lambda carry, x: module(carry, x), initial_carry, xs)`. Alternatively, consider using `equinox.filter_scan` if available, which is designed to handle Equinox modules.