Fold the AffineTransform bias into the matmul instead of concatenating a ones column

[cboulay]

Stacked A|B weights express y = xA + B. The implementation glued a column of ones onto the message and did one matmul, which copies the whole message every cycle just to carry a constant:

data = xp.concat((data, xp.ones(sample_shape)), axis=axis_idx)
data = xp.matmul(data, self._state.weights)

This replaces that with addmm where the backend has it, and an in-place bias add where it doesn't.

Correction: the first revision of this PR claimed broadcasting the bias is "cheaper on every backend." That was never measured — only MLX was — and it is false. The NumPy numbers below are the corrected result, and the fix that makes the non-MLX path actually worth taking is the in-place add. Thanks to @cboulay for catching it.

MLX — addmm

Folds the add into the matmul epilogue. Jittered 30–64 sample chunks, min of 5, A/B interleaved:

	before	after
256 ch, stacked A|B	26.2–26.9 µs/msg	23.1–24.1	~1.13x
1024 ch, stacked A|B	36.2–36.9 µs/msg	26.7–26.9	~1.37x
plain A (control)	25.0–25.8	24.8–26.8	unchanged

NumPy — the broadcast alone is a wash; the in-place add earns it

Dropping the concat saves copying the message but adds a full read-modify-write pass over the result. Those nearly cancel: matmul(...) + bias measured 0.80–1.15x against concat from 30x64 to 3000x1024, float32 and float64.

Adding in place skips the second allocation that + bias would immediately discard. Interleaved, min of 15 rounds, ratio vs concat:

n_ch	n_t=30	300	600	2000	6000
64	0.96x	0.96x	0.98x	0.97x	1.34x
256	0.98x	0.93x	1.04x	1.59x	1.57x
1024	0.96x	1.28x	1.20x	1.13x	1.24x

2–10% slower while the message still fits in cache, 1.1–1.7x faster once it doesn't. Taken unconditionally because the sides are asymmetric in absolute terms — +0.13 µs at 30x256 against −233 µs at 3000x256 — not because it wins everywhere. A size gate would win at every shape, but costs a machine-dependent constant to maintain (as util/blockdiag.py already does) for an effect capped at ~10% on the losing side.

A hypothesis tested and rejected

The concat penalty is not n_ch + 1 making the matmul's inner dimension odd. At n_ch=255 — so K=256 exactly, fully aligned — concat is still 1.53x slower at 3000 samples, and the effect is flat across 255/256/257 and 511/512/513. It's the copy's memory bandwidth, which is why the crossover tracks working-set size rather than shape. Recorded so nobody re-runs it.

On "in-place"

It mutates only the buffer matmul just allocated, which nothing else holds a reference to. The caller's message is never touched — which matters here, since other subscribers to the same message would otherwise silently see transformed data. Tests assert it for float32, float64 and int32 messages, that the result aliases no input, and that processing the same message twice is stable.

bias is a row of the same weight matrix as weights, so its dtype can never be wider than the matmul result's and the cast is always safe; the try/except covers a backend with immutable arrays (MLX rebinds rather than mutates there, and takes addmm anyway).

Structure

The A|B split is cached as two views on the stored weight matrix — free to keep, invalidated at all three sites where the weights are replaced. Hoisting the axis permute above the branch lets the stacked and plain paths share it.

The grouped-rereference data - (data @ project) @ spread was measured as an addmm(alpha=-1) candidate too and is not included: 0.90x–1.31x across repeats, never reproducibly better.

Tests

4068 pass, including two new regression tests for the non-mutation property above.