eqx-zoo 0.3.3

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

Decoder-based embedding models such as Qwen3-Embedding work the same way, through DecoderEmbedder. For search, queries take the instruction prompt from the model's sentence-transformers configuration; documents are embedded as they are.

from eqx_zoo import DecoderEmbedder

model = DecoderEmbedder.from_pretrained("Qwen/Qwen3-Embedding-0.6B")
tokenizer = Tokenizer.from_pretrained("Qwen/Qwen3-Embedding-0.6B")
tokenizer.enable_padding()

prompt = (
    "Instruct: Given a web search query, retrieve relevant passages that answer the query\nQuery:"
)
batch = tokenizer.encode_batch(
    [prompt + "What is the capital of France?", "Paris is the capital of France."]
)
ids = jnp.array([e.ids for e in batch])
mask = jnp.array([e.attention_mask for e in batch])
query, passage = jax.vmap(model.embed)(ids, mask)
print(query @ passage)  # cosine similarity

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, and collected on the Hugging Face Hub in Verified in eqx-zoo.

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)

RobertaModel shares the XLM-RoBERTa implementation and is verified on tiny random models.

Decoder embedding models (DecoderEmbedder)

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 near-tied decisions flip even in Hugging Face's own bf16. Instead, on a 512-token input, bf16 decisions must match float32's wherever float32 is confident (its top two logits differ by more than 8 times Hugging Face's own bf16 noise), both in a full forward pass and in cached decoding. Precision-critical blocks (attention's softmax and LayerNorm) are also tested directly in bf16 against the reference arithmetic.

Encoders and decoder embedders are tested layer by layer against Hugging Face, and their embeddings must match sentence-transformers for a padded batch of sentences, in float32. In bfloat16, their hidden states and embeddings must be about as accurate as the reference libraries' own bfloat16.

On GPUs and TPUs, JAX's default float32 matrix multiplications may use reduced precision internally (TF32 on recent NVIDIA GPUs, bfloat16 on TPUs), so float32 results can differ noticeably from a float32 reference. Request the most accurate setting, which is true float32 on GPUs and close to it on TPUs, with jax.config.update("jax_default_matmul_precision", "highest"), or with jax.default_matmul_precision("highest"): for a block of code. Setting the environment variable JAX_DEFAULT_MATMUL_PRECISION=highest does the same without changing any code, which is a quick way to check whether precision explains a difference. The test suite requests highest precision.

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.