eqx-zoo 0.3.1

eqx-zoo provides pretrained models as plain Equinox modules, loaded straight from Hugging Face checkpoints and numerically verified against their reference implementations: language models against 🤗 Transformers, and sentence-embedding models against sentence-transformers. Since Transformers v5 removed its JAX models, eqx-zoo is a verified way to keep using them in JAX.

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. ...

Batched generation

Prompts of different lengths are left-padded and generated together. With greedy decoding, each row is identical to generating that prompt on its own.

from eqx_zoo import generate_batch

tokenizer.enable_padding(direction="left")
prompts = ["The capital of France is", "The largest planet in the solar system is"]
batch = tokenizer.encode_batch(prompts)

ids = jnp.array([e.ids for e in batch])
mask = jnp.array([e.attention_mask for e in batch])
tokens = generate_batch(model, ids, mask, max_new_tokens=20)

Sentence embeddings

Encoder.embed reads each checkpoint's sentence-transformers configuration and applies its own pooling and normalisation, so embeddings match sentence-transformers.

import jax
import jax.numpy as jnp
from tokenizers import Tokenizer

from eqx_zoo import Encoder

tokenizer = Tokenizer.from_pretrained("sentence-transformers/all-MiniLM-L6-v2")
tokenizer.enable_padding()
model = Encoder.from_pretrained("sentence-transformers/all-MiniLM-L6-v2")

batch = tokenizer.encode_batch(["The cat sits on the mat.", "A feline rests on a rug."])
ids = jnp.array([e.ids for e in batch])
mask = jnp.array([e.attention_mask for e in batch])
embeddings = jax.vmap(model.embed)(ids, mask)
print(embeddings @ embeddings.T)  # cosine similarities, as the embeddings are normalised

Models

Any checkpoint with a supported architecture loads with CausalLM.from_pretrained or Encoder.from_pretrained, from the Hugging Face Hub or a local directory. These checkpoints are verified by the test suite.

Language models (CausalLM)

Mixture-of-experts models, including Qwen3-30B-A3B, use the same Qwen3MoeForCausalLM architecture; experts are computed with grouped matrix multiplications. generate_batch currently processes their prompts one after another (#23).

Encoders and embedding models (Encoder)

Verification

Language models are tested against Hugging Face activations in two tiers:

• float32: every layer's output must match, and greedy generation must reproduce the reference output token for token.
• bfloat16: logits must be about as accurate as Hugging Face's own bfloat16, measured against its float32 output. Exact greedy agreement isn't required in bf16, since it drifts even between Hugging Face's own bf16 and float32 runs.

Encoders are tested layer by layer against Hugging Face, and their embeddings must match sentence-transformers for a padded batch of sentences, in float32; bfloat16 verification for encoders is tracked in #28.

Every architecture is also tested on tiny randomly initialised models, which cover code paths that no single checkpoint exercises. See tests/.

Development

git clone https://github.com/xquantize/eqx-zoo && cd eqx-zoo
uv sync
uv run pytest -m "not checkpoint"   # fast: tiny random models, no downloads
uv run pytest                       # full: also downloads and verifies checkpoints
uv run python benchmarks/generation.py Qwen/Qwen3-0.6B   # load, compile and throughput

Contributions are welcome. See CONTRIBUTING.md for setup, conventions and how to add a model.

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.