Switch eqx.nn.Linear.__call__ to x @ weight.T to eliminate vmap+grad permutes (gated on GPU benchmark)

[mvanhorn]

Summary

This plan has two parts: the code change, and the benchmark evidence required for the PR to be acceptable to the maintainer given the GPU-flat counterevidence on the sibling issue #1196.

Code change

1. In equinox/nn/_linear.py, locate the line out = self.weight @ x inside Linear.__call__ (around line 97 of HEAD). Replace with:

   out = x @ self.weight.T

2. The bias add (out = out + self.bias) is unchanged because the bias is broadcast over the leading dimensions in both orderings.

3. The Identity and Embedding classes in the same file do not call weight @ x; they are untouched.

4. The jaxtyping annotation on the return type (Float[Array, " out"] for the non-vmapped path) is unchanged; the shape of the returned single-vector case is still (out,).

Tests

5. In tests/test_nn.py, add a numerical-equivalence test that compares the new Linear.__call__(x) output against x @ weight.T + bias (the explicit form) and against the previous weight @ x + bias form, asserting both are jnp.allclose to within float32 ULPs for several shapes (vector, 2D vmap, 3D vmap).

6. Add an HLO smoke test using jax.jit(jax.vmap(linear))(x).lower(x).compile().runtime_executable() (or lower(...).as_text()) and assert that the lowered HLO for the vmapped path contains fewer transpose ops after the change than before. Skip on JAX versions where the HLO text format differs.

Benchmark evidence (mandatory for PR)

The PR body MUST include:

• A reproducible benchmark script run on both CPU and at least one GPU (Colab T4 or local CUDA).
• The 3D-vmap residual MLP from the issue body, exercised at the listed shapes (DIM=64, BATCH=256, SEQ=32) and at one larger production-shaped configuration.
• A side-by-side table: w @ x vs x @ w.T for forward-only, forward+backward, and vmap(grad(...)).
• HLO transpose counts before/after for the vmapped path.
• An honest "GPU is approximately breakeven, CPU is materially faster, no regression observed" framing, NOT a claim of 3.5x speedup (that number is CPU-only).

7. If the GPU benchmark on the reviewer's hardware (the user can ask for a Colab link or run it themselves) shows any regression on a representative shape, do not push the PR - leave a benchmark-data comment on the issue instead and close out the plan. This is a salvage path noted in front-matter.

PR title: Use x @ weight.T in eqx.nn.Linear to eliminate redundant permutes under vmap+grad

PR body skeleton (concise, factual; fill the benchmark table from a real run):

Closes #1195.

Current Linear.__call__ does weight @ x. Under vmap, this puts the output dim before the batch dim, so XLA emits a transpose at every layer boundary to restore batch-leading layout. Switching to x @ weight.T matches Flax's Dense / PyTorch's Linear ordering and removes those transposes from the lowered HLO.

The single-vector path is numerically identical (verified in tests). Existing Linear-using modules (MLP, MultiheadAttention, LayerNorm's scale/bias path, etc.) are unaffected at the API level.

Benchmarks (script attached as a gist for reproducibility):

Device	Shape	w @ x	x @ w.T
CPU (M3)	8 blocks, (256, 32), DIM=64
GPU (T4)	same
GPU (T4)	larger production shape

HLO transpose count for the vmapped path drops from to .

Why this matters

eqx.nn.Linear.__call__ is implemented as out = self.weight @ x. For a single vector x of shape (in,) this is a plain matvec. But under jax.vmap, the batch dimension lands after the output dimension of the result, producing intermediates with shape (out, batch, ...) instead of (batch, ..., out). JAX/XLA then inserts a transpose at every layer boundary to restore batch-leading layout, because the surrounding loss / norm / activation / next-Linear stack all expects batch-leading.

The reporter's HLO analysis shows the transpose ops are real, not folded away. On CPU + double-vmap (e.g. an 8-layer MLP with (batch=256, seq=32)), this measures as a 3.5x slowdown vs. an equivalent module using x @ weight.T. On GPU the gap is much smaller (ZagButNoZig measured 0.628ms vs 0.501ms - ~25% on a tiny model). The fix is to write the matmul in the form that keeps batch dimensions leading: out = x @ self.weight.T (the convention used by Flax Dense and PyTorch Linear.__call__).

For the un-vmapped path (x is a single vector of shape (in,)), weight @ x and x @ weight.T produce bit-identical results, so this is a pure-perf change.

Testing

• uv run pytest tests/test_nn.py -k linear passes; new numerical-equivalence test confirms jnp.allclose for the single-vector, single-vmap, and double-vmap cases at fp32 and fp64.
• uv run pytest full suite passes; in particular, downstream tests for MLP, MultiheadAttention, Sequential containing Linear all stay green (no API or numerical change at the public surface for un-vmapped calls).
• New HLO-count test confirms strictly fewer transpose ops in the vmapped lowered IR after the change.
• Manual CPU benchmark from the issue reproducer (gist eqx_linear_vmap_3d-py) is reproduced locally; numbers go into the PR body.
• Manual GPU benchmark on Colab T4 (or local CUDA) confirms NO regression at any tested shape. If a regression appears, abort the PR per the plan's salvage path and post the data on the issue thread.
• uv run pyright equinox/nn/_linear.py passes; the change is a single rewrite of an expression and does not alter any type annotation.
• git diff --stat shows ~2 lines changed in equinox/nn/_linear.py plus the new test block in tests/test_nn.py.

Fixes #1195

AI was used for assistance.

[patrick-kidger]

Thanks for this! Unfortunately it looks like your new test is failing in CI?

[mvanhorn]

Thanks for catching this @patrick-kidger. Two things:

1. Two of the three failures (test_linear_matmul_order_equivalence[shape1/shape2]) were a bug in MY new test: explicit + linear.bias with shapes (5,4)+(4,) and (2,5,4)+(4,) is rejected by jax_numpy_rank_promotion='raise'. Fixed in 95f62bf by broadcasting bias to explicit.shape before the add.

2. The third failure (test_make_jaxpr::test_basic) is more interesting and I think a legitimate concern about the change itself. Before this PR, the unvmapped MLP path had zero transposes in the jaxpr. After, it has 9. That's the opposite of what this PR is supposed to do for that code path - the optimization targets vmap+grad, but for plain MLP forward passes x @ weight.T adds an explicit transpose where weight @ x didn't need one.

Whether those 9 transposes survive XLA optimization on real hardware is exactly the empirical question your GPU-benchmark gate is for. I don't have GPU access to run the benchmark side of this PR; the CPU side I can run but the issue body specifically called out CPU-only results are misleading.

If you'd rather I close this and post the CPU benchmark on #1195 as data instead of as a PR, happy to do that. Otherwise, suggestions on how you'd want to handle the test_basic expected-primitive count (update the constant, or treat the regression as blocking)?

[patrick-kidger]

No strong feelings on my part. Ideally these distinctions shouldn't actually matter in JAX-level code, and would all be washed away by XLA.

It might be that it's possible to resolve this by adding a new algebraic pattern match in to XLA? In fact I think ideally there would be two:

1. if one writes out weight @ x and this is vmap'd, then the pattern match could substitute in the faster form.
2. if one writes out x @ weight.T and this is not vmap'd, then a different pattern match could substitute in the other faster form.

WDYT?

[mvanhorn]

I think the XLA pattern match is the right long-term home for this - it fixes the pattern for everyone writing either form, not just equinox users, and it would also wash away the 9-transpose regression I measured on the unvmapped path here.

On the two matches: agreed they're duals. In jaxpr terms both forms lower to dot_general plus (sometimes) an explicit transpose, so the match would be on dot_general dimension_numbers rather than a literal transpose op, which XLA's algebraic simplifier seems well-placed to handle. I can file an XLA issue with the vmap'd/unvmapped jaxprs and the GPU benchmark from this PR as the motivating case, and link it back here.

For this PR meanwhile, happy either way: park it pending the XLA discussion, or keep the change behind the benchmark gate as a near-term win for the vmap+grad path. Your call.

[patrick-kidger]

Okay! I think let's park this one for now and see how the XLA discussion goes. If it goes well then we can keep our code in its current simple form!