eqx-zoo 0.1.0

eqx-zoo provides pretrained models as plain Equinox modules, loaded straight from Hugging Face checkpoints and numerically verified against the reference implementation in 🤗 Transformers.

Every model is an ordinary pytree, so jax.jit, jax.grad, jax.vmap and the rest of the JAX ecosystem work on it directly.

Installation

pip install eqx-zoo

Requires Python 3.12+. The example below also uses pip install tokenizers.

Quick example

import jax.numpy as jnp
from tokenizers import Tokenizer

from eqx_zoo import CausalLM, generate

tokenizer = Tokenizer.from_pretrained("Qwen/Qwen3-0.6B")
model = CausalLM.from_pretrained("Qwen/Qwen3-0.6B")

prompt = jnp.array(tokenizer.encode("The capital of France is").ids)
tokens = generate(model, prompt, max_new_tokens=30)
print(tokenizer.decode(tokens.tolist()))
# Paris. The capital of Italy is Rome. The capital of Spain is Madrid. ...

Models

Any checkpoint with a supported architecture loads with CausalLM.from_pretrained. These checkpoints are verified by the test suite:

Each verified checkpoint is tested against Hugging Face activations captured in float32: every layer's output must match, and greedy generation must reproduce the reference output token for token. See tests/.

Development

git clone https://github.com/xquantize/eqx-zoo && cd eqx-zoo
uv sync
uv run python scripts/make_reference.py   # capture reference activations
uv run pytest

Contributions are welcome. AGENTS.md describes the conventions every model follows.

See also

• Neural networks: Equinox.
• Typing: jaxtyping.
• Other JAX model collections: Levanter (LLM training), Bonsai (Flax NNX), Equimo (vision).

License

Apache 2.0. Pretrained weights are distributed by their original authors under their own licenses.