`nn.Upsample(mode="linear")` is 2-4x slower than a matmul-based bilinear resize when upsampling

Description

mx.nn.Upsample with mode="linear"/align_corners=False is significantly
slower than it needs to be specifically when upsampling (scale_factor > 1).
I found this while porting a video frame-interpolation model (RIFE) to MLX,
where repeated bilinear resizes are a meaningful fraction of the forward pass.

A fixed-ratio bilinear resize is separable into two 1D linear resizes, each of
which is just a matrix multiply by a small, fixed (2-nonzeros-per-row) weight
matrix — turning "resize" into "matmul" plays to MLX's own strength here. This
is bit-exact with nn.Upsample's output (see repro) and 2-4x faster for
upsampling shapes; for downsampling nn.Upsample is actually a bit faster
than the matmul version, so this looks like a real, specific gap in the
upsampling path rather than a case for just always avoiding nn.Upsample.

Repro

"""Repro: mx.nn.Upsample(mode="linear") is 4-8x slower than a
fixed-ratio bilinear resize implemented as two small matmuls, at
ordinary image-pipeline resize shapes. No torch involved - MLX only.
"""
import time
import numpy as np
import mlx.core as mx
import mlx.nn as nn


def matmul_resize_weight(in_size, out_size):
    scale = in_size / out_size
    src = (np.arange(out_size) + 0.5) * scale - 0.5
    src = np.clip(src, 0, in_size - 1)
    x0 = np.floor(src).astype(np.int64)
    x1 = np.minimum(x0 + 1, in_size - 1)
    w1 = (src - x0).astype(np.float32)
    w0 = 1.0 - w1
    w = np.zeros((out_size, in_size), dtype=np.float32)
    w[np.arange(out_size), x0] += w0
    w[np.arange(out_size), x1] += w1
    return mx.array(w)


def matmul_resize(x, scale_factor):
    n, h, w, c = x.shape
    out_h, out_w = round(h * scale_factor), round(w * scale_factor)
    wh, ww = matmul_resize_weight(h, out_h), matmul_resize_weight(w, out_w)
    x = (wh @ x.transpose(1, 0, 2, 3).reshape(h, n * w * c)).reshape(out_h, n, w, c).transpose(1, 0, 2, 3)
    x = (ww @ x.transpose(2, 0, 1, 3).reshape(w, n * out_h * c)).reshape(out_w, n, out_h, c).transpose(1, 2, 0, 3)
    return x


def timed(fn, n=30, warmup=8):
    for _ in range(warmup):
        mx.eval(fn())
    t0 = time.perf_counter()
    for _ in range(n):
        mx.eval(fn())
    return (time.perf_counter() - t0) / n * 1000


for in_h, in_w, scale in [(192, 320, 4.0), (384, 640, 2.0), (768, 1280, 0.25), (768, 1280, 0.5)]:
    x = mx.random.normal((1, in_h, in_w, 4))
    up = nn.Upsample(scale_factor=scale, mode="linear", align_corners=False)
    up_c = mx.compile(lambda a: up(a))
    resize_c = mx.compile(lambda a: matmul_resize(a, scale))

    up_ms = timed(lambda: up_c(x))
    resize_ms = timed(lambda: resize_c(x))
    diff = float(mx.abs(up_c(x) - resize_c(x)).max())
    out_h, out_w = round(in_h * scale), round(in_w * scale)
    print(f"{in_h}x{in_w} -> {out_h}x{out_w} (scale={scale}): "
          f"nn.Upsample={up_ms:.3f}ms  matmul_resize={resize_ms:.3f}ms  "
          f"({up_ms/resize_ms:.1f}x slower)  max_diff={diff:.2e}")

Output (M1 Pro, mlx 0.31.2, macOS 26.4.1)

192x320 -> 768x1280 (scale=4.0): nn.Upsample=4.008ms  matmul_resize=0.979ms  (4.1x slower)  max_diff=4.77e-07
384x640 -> 768x1280 (scale=2.0): nn.Upsample=3.937ms  matmul_resize=1.886ms  (2.1x slower)  max_diff=4.77e-07
768x1280 -> 192x320 (scale=0.25): nn.Upsample=0.509ms  matmul_resize=0.863ms  (0.6x slower)
768x1280 -> 384x640 (scale=0.5): nn.Upsample=1.180ms  matmul_resize=1.659ms  (0.7x slower)

(Both compiled with mx.compile; same pattern holds eager, just uniformly
slower on both sides.)

Context

This showed up as a real end-to-end cost: in a RIFE (video frame
interpolation) port, replacing nn.Upsample with this matmul-based resize
cut ~18ms off a ~354ms forward pass at 768x1280 purely from the upsample
calls (each block does one upsample of its flow/mask output back to full
resolution). Happy to open a PR with the matmul-based implementation
(resize()/_weight_matrix() in the repro) if that's a fix you'd want,
or this may point to something fixable at a lower level in however
upsample_linear currently walks the output for scale_factor > 1.

For context on where MLX is already doing well on this same machine (so
this issue is a targeted gap, not a general performance complaint): a
native mx.nn/mx.core GPT (minGPT's gpt-mini config, 2.71M params),
compiled with mx.compile's stateful pattern
(mx.compile(step, inputs=state, outputs=state)), trains 1.76x faster
than the same architecture on PyTorch+MPS (88.7ms/step vs. 154.9ms/step)

• consistent with transformer_lm's existing example using the same
  compile pattern. The Upsample gap above is the one place in two separate
  model ports (this GPT, and the RIFE port) where MLX measured slower than
  PyTorch on identical work.

Environment

• mlx 0.31.2
• macOS 26.4.1, Apple M1 Pro

Start at the nn.Upsample entry point and the upsample_linear path mentioned in the issue; use the supplied benchmark to compare scale_factor > 1 against downsampling and verify numerical equivalence. Done means the upsampling path closes the reported performance gap without changing the output.