Small input linear and cubic upsampling should use dense matmuls

[slink]

Gather-based interpolation backward is a scatter-add, which is slow on Metal. For upscaling with small spatial inputs, we can instead build per-axis interpolation matrices and contract them.

I found this to be 2-4x faster forward and backward for linear with N <= 128 and cubic gated with N <= 32.

• ☑️ I understand it is strictly prohibited to use AI to write PR description
• AI usage disclosure: I used AI to help me benchmark this change

Benchmark data

Setup:

• Apple M5 Max (MacBook Pro)
• macOS 26
• mlx 0.32.2
• float32
• MLX_ENABLE_TF32=0

Times in ms, mean of 10–30 calls after warmup; "fwd+bwd" is mx.grad of a mean-square loss through the op.

crossover sweeps

Crossover sweep — linear, ×2 upscale, input (4, N, N, 32)

N_in	matmul fwd	gather fwd	matmul fwd+bwd	gather fwd+bwd	gated path
8	0.18	0.17	0.20	0.22	matmul
16	0.17	0.28	0.20	0.39	matmul
32	0.21	0.52	0.18	0.51	matmul
64	0.41	0.76	0.33	1.06	matmul
96	0.77	1.54	0.65	2.38	matmul
128	1.14	2.63	1.08	4.48	matmul
192	2.59	6.18	2.80	9.55	gather
256	4.27	11.92	5.77	17.06	gather
384	12.59	25.95	16.84	39.52	gather
512	21.60	893.76	92.01	71.28	gather

Cases the gate excludes (gather wins or matmul loses)

case	matmul fwd	gather fwd	matmul fwd+bwd	gather fwd+bwd
N=256, ×0.5 downscale	1.14	0.73	0.88	2.06
N=512, ×0.5 downscale	20.38	14.50	5.11	7.73
N=512, ×2 (backward)	—	—	92.01	71.28

Other measured points

case	matmul	gather
linear (64, 8, 16, 64) ×10, fwd	4.60	16.81
linear (64, 8, 16, 64) ×10, bwd	3.90	27.57
cubic (4, 32, 32, 32) ×2, fwd	0.40	0.93

Testing

• Full python/tests/test_upsample.py passes: 9 tests, 172 subtests (was 6 / 152 before the PR), including torch parity across fractional scales and both align_corners settings.
• Matmul path agrees with the gather path to 2.4e-7 max deviation (with MLX_ENABLE_TF32=0; see caveats).
• New tests added in the PR:
  • test_matmul_gather_path_parity — matmul vs gather output equality across cases straddling both thresholds, including a downscale case, a 1D case, and both align_corners settings.
  • test_matmul_gather_grad_parity — gradient equality at the motivating shape (8×16, ×10), since the backward pass is the point of the change.
  • test_matmul_path_routing — exact threshold boundaries for both modes (tied to the constants, so retuning keeps the test valid), mixed up/down scale rejected, scale=1.0 accepted.
• The new tests error on main (_interpolate_matmul and the threshold constants do not exist there), demonstrating they exercise only new code.
• pre-commit (black, isort) clean on both changed files.

Caveats

• Thresholds are measured on one machine (M5 Max).
• Cubic is gated at 32 on a single measured point; its crossover is lower than the linear one, but it has not been swept.
• TF32 interaction: the matmul path runs through GEMM, so under the default MLX_ENABLE_TF32=1 on M5 Metal / CUDA it loses ~1e-3 accuracy relative to the gather path (which never touches GEMM). With MLX_ENABLE_TF32=0, parity is 2.4e-7. Relates to [BUG] fp32 matmul silently defaults to TF32-class precision (MLX_ENABLE_TF32=1), undocumented on both backends #3860.