Eqx-zoo

eqx-zoo 0.3.4 verifies gradients against PyTorch, and fixes two training-behaviour differences that the new checks found.

Highlights

- **Gradients verified.** For every model class (causal LMs including mixture-of-experts, encoders and decoder embedders), each parameter's gradient matches PyTorch's to within about 1e-6, on randomly initialised test models. These checks catch bugs that leave every forward value unchanged, such as a stray `stop_gradient`.
- **GPU support on Linux.** `pip install "eqx-zoo[cuda]"` installs JAX's CUDA build on Linux with an NVIDIA GPU (including WSL2); on macOS and native Windows the extra changes nothing. Tested on an NVIDIA GPU.

Fixes

- **Padding embeddings are no longer trained**, matching PyTorch's `padding_idx`. Previously, gradients reached the padding token's embedding row: in encoders (word embeddings, and RoBERTa's padding position) and in decoders with a pad id. Forward results are unchanged, but fine-tuning now matches Hugging Face. Tied output heads still train that row, as in PyTorch.

API changes

- **New:** `Config.pad_token_id`, read from the checkpoint's `config.json`.

Project

- Plain `pytest` now runs only the fast tests, so it's safe on any laptop; checkpoint tests run with `--checkpoint NAME`, or all of them with `scripts/test-checkpoints.sh`, one process each (about 7 GB peak). Every run reports its peak memory.
- The bf16 attention test is robust across CPUs and GPUs, and the long-input decision tests are calibrated on CI's x86 runners. They're skipped for mixture-of-experts models, whose routing choices can flip between platforms.
- Pull requests now need an approving review.

Thanks to SamuelWLee27 for testing on an NVIDIA GPU and for the GPU fixes.

Known issues

- JIT-compiled bf16 is less accurate than eager bf16 for causal LMs on CPU; excess precision explains part of it (11).
- `generate_batch` processes prompts for MoE models one after another (23).

Install

pip install --upgrade eqx-zoo
pip install --upgrade "eqx-zoo[cuda]" Linux with an NVIDIA GPU

eqx-zoo 0.3.3 adds decoder-based embedding models and strengthens numerical verification, much of it prompted by a detailed community report from an NVIDIA L4.

Highlights

- **Qwen3-Embedding.** The new `DecoderEmbedder` loads decoder embedding models such as [Qwen3-Embedding-0.6B](https://huggingface.co/Qwen/Qwen3-Embedding-0.6B), with last-token pooling and normalisation read from the checkpoint's sentence-transformers configuration. Embeddings match sentence-transformers in float32 (max difference 4.7e-7, including with the query instruction prompt) and are verified in bf16.
- **Decision-level bf16 tests.** On a 512-token input, bf16 decisions must match float32's wherever float32 is confident, both in a full forward pass and in cached decoding. These tests catch position bugs that only appear beyond 256 tokens, which the short-prompt tests miss.
- **Block-level bf16 tests.** Attention's softmax and LayerNorm are tested directly in bf16 against the reference arithmetic, catching precision bugs too subtle to see at the model level.
- **True float32 on accelerators.** The test suite requests JAX's highest matmul precision, and the README explains why float32 results on GPUs and TPUs depend on it.

API changes

- **New:** `DecoderEmbedder` and `eqx_zoo.layers.last_token_pool`.
- `Config.from_hf` recognises base-model architecture names (`LlamaModel`, `Qwen2Model`, `Qwen3Model`), as used by headless decoders.

Project

- An issue template for reporting results on GPUs and TPUs, linked from `CONTRIBUTING.md`.
- The verified checkpoints are collected on the Hugging Face Hub: [Verified in eqx-zoo](https://huggingface.co/collections/xquantize/verified-in-eqx-zoo-6ac2d731fd1ef3e109541bf7).

Thanks to the author of the L4 report on the Hugging Face forum, whose findings led directly to the decision-level and precision changes in this release.

Known issues

- JIT-compiled bf16 is less accurate than eager bf16 for causal LMs on CPU; excess precision explains part of it (11).
- `generate_batch` processes prompts for MoE models one after another (23).
- Scan-over-layers could speed up compilation on accelerators (15).

Install

pip install --upgrade eqx-zoo

eqx-zoo 0.3.2 adds multilingual embedding models and verifies encoders in bfloat16.

Highlights

- **RoBERTa and XLM-RoBERTa.** `Encoder.from_pretrained` now loads `RobertaModel` and `XLMRobertaModel` checkpoints. Together with BERT, the supported encoder architectures account for about 81% of embedding-model downloads on the Hugging Face Hub.
- **Multilingual embeddings.** [multilingual-e5-base](https://huggingface.co/intfloat/multilingual-e5-base) (XLM-RoBERTa) and [multilingual-e5-small](https://huggingface.co/intfloat/multilingual-e5-small) (BERT) are verified: they match Hugging Face layer by layer, and their embeddings match sentence-transformers.
- **Encoders verified in bfloat16.** In bf16, encoder hidden states and embeddings must be about as accurate as Hugging Face's and sentence-transformers' own bf16, in both eager and JIT-compiled execution. `LayerNorm` is also tested directly against PyTorch in bf16, which catches subtle precision bugs that are too small to see at the model level.

API changes

- **New:** `EncoderConfig.pad_token_id` and `EncoderConfig.padding_aware_positions`, with defaults that keep BERT's behaviour.

Project

- CI runs one job per checkpoint, in parallel, each with its own cache. This fixes out-of-memory failures, and adding a model no longer invalidates every other model's cache.
- Checkpoint tests are labelled with the checkpoint they need; `pytest --checkpoint NAME` runs one checkpoint's tests.

Known issues

- `generate_batch` processes prompts for MoE models one after another (23).
- JIT-compiled bf16 is less accurate than eager bf16 for causal LMs (11).
- The bf16 tests don't cover long-context softmax precision (12).

Install

pip install --upgrade eqx-zoo

eqx-zoo 0.3.1 adds sentence-embedding models: BERT-style encoders, verified against sentence-transformers.

Highlights

- **Encoders.** `Encoder.from_pretrained` loads `BertModel` checkpoints, including sentence-transformers models such as [all-MiniLM-L6-v2](https://huggingface.co/sentence-transformers/all-MiniLM-L6-v2) and [bge-small-en-v1.5](https://huggingface.co/BAAI/bge-small-en-v1.5).
- **Sentence embeddings.** `Encoder.embed` reads each checkpoint's sentence-transformers configuration and applies its own pooling (mean or `[CLS]`) and normalisation. Embeddings match sentence-transformers for padded batches of sentences.
- **Verified.** Both checkpoints and a tiny random BERT match Hugging Face layer by layer in float32.
- **New building blocks** in `eqx_zoo.layers`: `LayerNorm`, `dot_product_attention`, `mean_pool`, `cls_pool` and `l2_normalize`.

Project

- Model classes now live in `eqx_zoo.models`. The public API is unchanged: `from eqx_zoo import CausalLM, Config, Encoder`.
- `CONTRIBUTING.md` documents the versioning policy: patch releases for backwards-compatible changes, minor releases only for breaking changes to the public API.

Known issues

- Encoders are verified in float32; bfloat16 verification is tracked in 28.
- `generate_batch` processes prompts for MoE models one after another (23).
- JIT-compiled bf16 is less accurate than eager bf16 for causal LMs (11).

Install

pip install --upgrade eqx-zoo

eqx-zoo 0.3 adds the first mixture-of-experts architecture.

Highlights

- **Qwen3-MoE.** `CausalLM.from_pretrained` loads `Qwen3MoeForCausalLM` checkpoints, the architecture of Qwen3-30B-A3B. Each token is routed to its top-k experts, and experts are computed together with grouped matrix multiplications (`jax.lax.ragged_dot`), so compute scales with the active experts rather than all of them.
- **Verified on a trained model.** [MiniMind-3 MoE](https://huggingface.co/jingyaogong/minimind-3-moe) passes the full parity suite: layer by layer, logits, greedy generation, KV cache, bf16, padding and batching. A tiny random-weight Qwen3-MoE covers top-2 routing, mixed MoE/dense layers and an untied head.
- **Mixed layers.** Models can combine MoE and dense layers, following `decoder_sparse_step` and `mlp_only_layers` as in the reference implementation.
- **Both config formats.** The expert count is read from `num_experts` (transformers 4.x) or `num_local_experts` (5.x).

API changes

- **New:** `eqx_zoo.layers.SparseMoE`. `Config` gains `num_experts`, `num_experts_per_tok`, `moe_intermediate_size`, `norm_topk_prob` and `moe_layers`, with defaults meaning a dense model. `DecoderLayer` takes a `moe` flag.

Project

- `CONTRIBUTING.md` describes setup, conventions, how to add a model, and the release process.
- `CITATION.cff` enables GitHub's "Cite this repository".
- Dependabot keeps the CI actions up to date.

Known issues

- `generate_batch` processes prompts for MoE models one after another, because `jax.lax.ragged_dot` cannot yet be vmapped (23). Results are exact.
- JIT-compiled bf16 is less accurate than eager bf16 (11).
- The bf16 tests don't cover long-context softmax precision (12).

Install

pip install --upgrade eqx-zoo

Fixes the README on PyPI: the banner and links now use absolute URLs, so they render on the project page. No code changes since v0.2.0.