Equinox

Equinox is a JAX library for building models as PyTrees, applying filtered JAX transformations, and using reusable neural-network layers, stateful modules, and advanced debugging utilities.

Use this skill when the user mentions:

• equinox, eqx, eqx.Module, or equinox.nn
• filtered transforms such as filter_jit, filter_grad, filter_vmap, or filter_pmap
• model surgery, PyTree partitioning, serialization, or runtime errors
• Sequential, MLP, BatchNorm, State, Shared, or inference_mode
• equinox.internal helpers such as noinline, while_loop, or scan

Quick start

Install the public package with:

pip install equinox

For a local checkout or a health check, run the bundled scripts/smoke.py from the generated skill root (the directory containing this SKILL.md). Set EQUINOX_SKILL_ROOT to that absolute directory and keep it as the command working directory:

cd "$EQUINOX_SKILL_ROOT" && python scripts/smoke.py --mode all

If you need multi-device CPU coverage for pmap or sharding checks, add --two-cpu-devices to the smoke script.

Route map

sub-skills/module-and-trees

Use for eqx.Module, field declarations, abstract attributes, PyTree filtering, tree_at, tree_equal, tree_check, apply_updates, Partial, and model surgery.

sub-skills/filtered-transformations

Use for filter_jit, filter_grad, filter_vmap, filter_pmap, filter_eval_shape, filter_make_jaxpr, custom filtered AD, checkpointing, callbacks, and sharding behavior.

sub-skills/nn-and-state

Use for equinox.nn layers, Sequential, Lambda, MLP, BatchNorm, State, StateIndex, Shared, and inference-mode workflows.

sub-skills/diagnostics-and-serialization

Use for tree_serialise_leaves, tree_deserialise_leaves, runtime errors, debug tools, pretty printing, enumerations, progress meters, and package troubleshooting.

sub-skills/internal-advanced

Use only for advanced or downstream-library work with equinox.internal, including noinline, while_loop, scan, nontraceable, finalise_jaxpr, and primitive authoring helpers.

How to choose

• If the task is about model shape, PyTree composition, or field semantics, start with module-and-trees.
• If the task is about JAX transforms crossing model boundaries, start with filtered-transformations.
• If the task is about built-in layers, training loops, or stateful modules, start with nn-and-state.
• If the task is about debugging, serialization, or runtime checks, start with diagnostics-and-serialization.
• If the task explicitly names equinox.internal or an advanced helper that is not part of the everyday public API, start with internal-advanced.

Reading order

1. Read the sub-skill that matches the dominant workflow.
2. Use the linked reference file in that sub-skill for API details, examples, and troubleshooting.
3. Use references/api-reference.md when you need a cross-skill API index.
4. Use references/troubleshooting.md when the issue looks like install/import, backend, optional dependency, or mixed-leaf JAX behavior. Read references/repo-provenance.md before relying on version-sensitive guidance.

Common signals

• A plain jax.jit or jax.grad call fails because a model contains non-array leaves: route to filtered-transformations.
• A module definition fails because fields or abstract attributes are missing: route to module-and-trees.
• A layer needs state or axis_name, or Sequential has a stateful layer: route to nn-and-state.
• A runtime error, serialization round-trip, or debug trace is the focus: route to diagnostics-and-serialization.
• A user asks about noinline, while_loop, scan, or other semi-public helpers: route to internal-advanced.

Notes

• Equinox is intentionally PyTree-first: models, helper wrappers, and advanced utilities all depend on ordinary JAX semantics rather than a separate runtime.
• There is no repo CLI. Treat library installation, import checks, and API inspection as the primary entry points.
• Prefer the bundled references over source checkout links so the skill stays usable after the repository is no longer present.