Fix matmul() broadcasting for operands of different ranks - #239

np.matmul() puts the batch dims in the wrong place when the operands have different ranks and the right one has batch dims (2D @ 3D+, or 3D+ @ 3D+ with different ranks). Sometimes the shape still looks right and only the values are wrong:

const a = np.array([[1, 2], [3, 4]]);
const b = np.array([[[1, 0], [0, 1]], [[2, 0], [0, 2]]]); // [I, 2I]
np.matmul(a, b).js();
// before: [[[1, 2], [2, 4]], [[3, 4], [6, 8]]]
// numpy:  [[[1, 2], [3, 4]], [[2, 4], [6, 8]]]  (that is, [a, 2a])

np.matmul(np.zeros([2, 3]), np.zeros([4, 3, 2])).shape;       // before [2, 4, 2], numpy (4, 2, 2)
np.matmul(np.zeros([7, 1, 2, 3]), np.zeros([4, 3, 5])).shape; // before [4, 7, 2, 5], numpy (7, 4, 2, 5)

NumPy and jax.numpy.matmul align batch dims from the right and broadcast them, which is also what np.einsum("...ij,...jk->...ik", a, b) already does here.

Cause

src/library/numpy.ts:1925 passes min(max(x.ndim, 2), y.ndim) - 2 trailing batch dims of each operand to lax.dot(). When the ranks differ, the leading dims that only the higher-rank operand has are passed as free dims, and lax.dot() returns [...batch, ...lhs free, ...rhs free], so they come out after the shared batch dims, or after the rows of x when they belong to y. This has been the case since #16 moved matmul() onto lax.dot().

Fix

Prepend 1s to the shape of the operand with fewer dims, so both have the same batch rank and lax.dot() broadcasts size-1 batch dims as it already does for equal ranks. A 1D x or a 2D y has no batch dims to align, so those paths are unchanged; for (..., N, K) @ (K, M), the usual dense-layer shape, the traced jaxpr is identical to before. MatMul in packages/onnx calls np.matmul() directly, so it picks up the fix.

Two tests are added to the jax.numpy.matmul() suite: a 2D matrix against a stack of matrices, and a 4D @ 3D product that broadcasts size-1 batch dims.

Verification

I could not install the dependencies (no npm registry access here), so vitest, pnpm check, pnpm lint and pnpm format:check were not run. Instead I loaded src/ directly in Node 22, with a small loader for extensionless imports and using, on the cpu and wasm backends:

• All rank pairs from 1D to 4D with batch sizes 1, 2 and 5 (150 shape pairs, integer-valued float32 data so results are exact), compared with NumPy 2.5.4, eagerly and under jit(): 97/150 match on main on each device, with every failure in 2D @ 3D/4D, 3D @ 4D or 4D @ 3D; 150/150 match with the fix.
• valueAndGrad() with respect to both operands for 9 mixed-rank cases, against analytic gradients computed in NumPy, eager and under jit(): 2/9 match on main, 9/9 with the fix. vmap() agrees with the unmapped call, both inputs are consumed exactly once, and incompatible batch dims still throw.
• I ran the tests in test/numpy.test.ts through a minimal stand-in for the vitest API. The two new tests fail on main on cpu and wasm (expected [2,3,2] to equal [3,2,2], expected [3,2,1,1] to equal [2,3,1,1]) and pass with the fix. The numpy, numpy-linalg, lax-linalg, einsum, refcount, basic and tracing test files give the same per-test results before and after apart from those four; the only failures, on both, are Float16Array is not defined from Node 22.
• JAX 0.12 returns the same shapes and values that both new tests expect.

WebGPU and WebGL were not available, so the new tests have not run there. The second one multiplies int32 arrays, as the existing einsum tests already do on every device.

🤖 Generated with Claude Code