Rewrite diag(A @ B) to (A * B.T).sum(-1) instead of materializing the product

[jessegrabowski]

There is no ExtractDiag rewrite for a matmul, so diag(A @ B) computes the full N×N product and indexes the diagonal out of it, where the equivalent (A * B.T).sum(-1) gets the same result in O(N·K) time and memory.

ExtractDiag already pushes through elementwise ops and has structural rewrites for block_diag/kron/alloc_diag (pytensor/tensor/rewriting/linalg/products.py). Matmul is an obvious next choice.

Microbenchmark

A shape (N, K), B shape (K, N), K = 10. pytest-benchmark median; async backends synced per call (block_until_ready for JAX, mx.eval for MLX).

• current: pytensor.function([A, B], pt.diag(A @ B), mode=...)
• proposed: pytensor.function([A, B], (A * B.T).sum(axis=-1), mode=...)

GPU backends (the interesting case — MLX on an Apple GPU, JAX on a Tesla T4):

MLX (Apple GPU, float32, pytensor 3.0.3)

N	current	proposed	speedup
10	117.0 µs	102.8 µs	1.1×
50	123.1 µs	111.0 µs	1.1×
100	122.8 µs	104.4 µs	1.2×
500	126.0 µs	105.0 µs	1.2×
2000	157.0 µs	107.7 µs	1.5×
8000	839.4 µs	121.7 µs	6.9×

JAX (Tesla T4, CUDA, pytensor 2.38.2)

N	current	proposed	speedup
10	143.4 µs	117.9 µs	1.2×
50	203.9 µs	119.7 µs	1.7×
100	210.2 µs	124.0 µs	1.7×
500	237.6 µs	120.8 µs	2.0×
2000	752.1 µs	151.1 µs	5.0×
8000	9.83 ms	336.1 µs	29.3×

CPU backends (C, NUMBA, JAX-CPU) — hundreds of × at large N

C — Apple Silicon

N	current	proposed	speedup
10	2.0 µs	2.1 µs	0.96× (slower)
50	3.3 µs	2.4 µs	1.4×
100	4.0 µs	2.8 µs	1.4×
500	114.9 µs	6.0 µs	19.2×
2000	336.0 µs	15.4 µs	21.9×
8000	25.62 ms	59.3 µs	432.4×

NUMBA — Apple Silicon

N	current	proposed	speedup
10	2.7 µs	2.7 µs	1.0×
50	4.3 µs	2.8 µs	1.5×
100	4.9 µs	3.0 µs	1.6×
500	114.1 µs	5.0 µs	22.8×
2000	433.6 µs	10.3 µs	42.1×
8000	26.63 ms	37.3 µs	714.1×

JAX (CPU) — Apple Silicon

N	current	proposed	speedup
10	23.0 µs	9.1 µs	2.5×
50	23.4 µs	8.8 µs	2.7×
100	27.2 µs	21.1 µs	1.3×
500	113.4 µs	22.5 µs	5.0×
2000	237.4 µs	29.1 µs	8.2×
8000	9.16 ms	38.5 µs	237.5×

C — Tesla T4 host (x86_64)

N	current	proposed	speedup
10	5.1 µs	5.3 µs	0.97× (slower)
50	6.2 µs	6.0 µs	1.0×
100	8.8 µs	6.9 µs	1.3×
500	24.9 µs	14.4 µs	1.7×
2000	576.6 µs	41.6 µs	13.8×
8000	104.10 ms	169.8 µs	613.3×

NUMBA — Tesla T4 host (x86_64)

N	current	proposed	speedup
10	7.1 µs	6.2 µs	1.1×
50	8.0 µs	6.8 µs	1.2×
100	10.2 µs	7.4 µs	1.4×
500	27.7 µs	11.6 µs	2.4×
2000	252.6 µs	27.9 µs	9.1×
8000	65.53 ms	118.0 µs	555.4×

Notes

• CPU: a large win once N ≳ 500 (hundreds of × at N = 8000), a wash below that, and marginally slower at N = 10 (0.96–0.97×) — the expected BLAS constant-factor crossover.
• GPU: parallelism makes the full matmul cheap, shrinking the advantage to ~7× (MLX) / ~29× (T4) at N = 8000; small-N timings are floored by fixed dispatch latency (~100–200 µs), equal for both variants. The matmul never wins — the reduction is faster or tied at every N on both GPUs.
• The rewrite should gate on the Dot having a single client (otherwise the product is computed anyway and the diagonal would be a redundant second pass), and given the small-N wash it may want a shape/size guard rather than firing unconditionally.