MLX: add native_group_norm and upsample_nearest2d handlers

[msluszniak]

Summary

Fixes #22017.

The MLX backend has no handler for aten.native_group_norm or aten.upsample_nearest2d. GroupNorm sits in every ResBlock of a Stable-Diffusion-style UNet and nearest upsampling sits in every decoder stage, so the two gaps together shatter a diffusion model rather than merely slowing it down.

Measured on the SDXS-512-DreamShaper UNet (SD-1.5 architecture, 4x64x64 latents):

	delegate subgraphs	nodes left on CPU
before	28	25x native_group_norm, 2x upsample_nearest2d
after	1	none

Each of those 27 boundaries was a delegate handoff per denoise call, leaving and re-entering the MLX runtime.

Both ops lower to primitives the backend already has, so this needs no schema or runtime change.

Approach

native_group_norm normalizes each group of C / group channels together with all of their spatial positions. Reshaping the input to (N * group, (C / group) * HxW) puts exactly that set on the last axis, so fast::layer_norm computes the normalization as a single fused kernel. The affine parameters are applied afterwards on the original shape rather than being passed to layer_norm, because group norm's weight and bias are per channel while layer_norm's are per normalized element; the two only coincide when every group holds a single channel. Only the normalized output is produced, matching the existing native_layer_norm handler's treatment of mean/rstd.

upsample_nearest2d becomes take(take(x, idx_h, -2), idx_w, -1). The source index for an output position is min(floor(dst * scale), in_size - 1) (aten's nearest_neighbor_compute_source_index), which depends only on the static input and output sizes, so both index vectors are constants. Expressing it as a gather rather than a repeat also covers non-integer scale factors and downsampling. Both the .vec and .default overloads are registered.

Test plan

Adds group_norm and upsample_nearest2d to backends/mlx/test/test_ops.py, 11 configurations in total:

• group norm: affine and non-affine, one channel per group (instance norm) and one group for all channels, a non-square spatial extent, and a 3D (N, C, L) input
• upsampling: integer, anisotropic and fractional scale factors, an explicit output size, and downsampling

All 11 match eager through the MLX runtime; the upsample ones are bit-exact (rtol = atol = 0).

python -m executorch.backends.mlx.test.run_all_tests group_norm upsample_nearest2d

cc @metascroy

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/22050

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

⚠️ 1 Awaiting Approval

As of commit bb67233 with merge base fbd4bbf ():

AWAITING APPROVAL - The following workflow needs approval before CI can run:

• periodic (gh)

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[linux-foundation-easycla]

The committers listed above are authorized under a signed CLA.

• ✅ login: msluszniak / name: Mateusz Słuszniak (bb67233)

[github-actions]

This PR needs a release notes: label

If your change should be included in the release notes (i.e. would users of this library care about this change?), please use a label starting with release notes:. This helps us keep track and include your important work in the next release notes.

To add a label, you can comment to pytorchbot, for example
@pytorchbot label "release notes: none"

For more information, see
https://github.com/pytorch/pytorch/wiki/PyTorch-AutoLabel-Bot#why-categorize-for-release-notes-and-how-does-it-work.

[metascroy]

@claude review this code

[claude]

Claude finished @metascroy's task in 6m 14s —— View job

---

Code Review

• Read changed files (ops.py, test_ops.py)
• Analyze native_group_norm handler
• Analyze upsample_nearest2d handler
• Review test coverage

Nice work — this is a clean, well-reasoned pair of handlers, and the commit message / PR description do an excellent job explaining why the decomposition is correct. Both handlers reuse existing primitives and follow the conventions of the surrounding code (native_layer_norm, the used_getitem_indices guard, require_args/require_kwargs, slot allocation). I ran through the logic carefully; below are the findings, ordered by importance.

---

1. Index computation precision vs. aten (potential correctness edge) — ops.py _nearest_source_indices

aten's nearest_neighbor_compute_source_index computes floor(dst_index * scale) in float32 (the kernel casts scale to float, see kernels/portable/cpu/util/upsample_util.h:140-149 and the const float scale_h signature in op_upsample_nearest2d.cpp:76). Eager PyTorch's reference does the same. This handler computes it in float64:

idx = (torch.arange(out_size, dtype=torch.float64) * step).to(torch.int64)

For the tested sizes this is bit-exact, but float32 and float64 can floor to different integers when dst * scale lands just below an integer boundary at larger output indices (e.g. tall/wide feature maps with a non-power-of-two scale — exactly the diffusion-decoder case this PR targets). Since the whole point of the backend is to match eager, matching aten's float32 arithmetic removes that risk:

step = ... # float
idx = (torch.arange(out_size, dtype=torch.float32) * torch.float32(step)).floor().to(torch.int64)

I'd flag this as worth confirming rather than a definite bug — the divergence only appears at boundary values the current tests don't hit. A test with a large odd spatial extent + fractional scale (e.g. shape=(1,3,127,127), scale_factor=(1.7,1.7)) would either confirm parity or surface it. Fix this →

2. native_group_norm reshape assumes NCHW-contiguous input

The reshape to (N*group, (C//group)*HxW) is only equivalent to per-group normalization if the group * (C//group) channels are laid out channel-major (standard NCHW contiguous), so that one group's channels are adjacent. This matches eager's own x.reshape(N, group, -1) decomposition and is almost certainly the layout the backend already works in, so this is likely fine — but it's an implicit assumption worth a one-line comment, since a channels-last input would silently produce wrong groupings rather than an error. Worth confirming the MLX partitioner never hands this handler a non-contiguous/channels-last tensor.

3. .vec detection via string match — minor robustness

if ".vec" in str(n.target):

This works, but comparing the target directly is less fragile and reads better:

if n.target is torch.ops.aten.upsample_nearest2d.vec:

---

Things I checked that look correct

• getitem guard: correctly rejects use of mean/rstd (indices 1/2), mirroring native_layer_norm. ✅
• Affine application: the four-way branch (no affine / weight-only / bias-only / both) routes the final write into the node's real output slot out in every path, with temporaries otherwise. Verified each branch. ✅
• affine_shape [1, C, 1, ...] broadcasts per-channel weight/bias correctly for both 4D and 3D (N,C,L) inputs (x_ndim - 2 trailing ones). ✅
• Static-shape guards: N/C/HxW/group and input/output spatial sizes are all validated as int with clear error messages. ✅
• Constant dedup: the constant name encodes in_size/out_size/scale, so identical index vectors are shared and distinct ones don't collide. ✅
• require_args bounds: native_group_norm (8,8) and both upsample overloads match the aten schemas. ✅

Test coverage

Coverage is thorough for the common cases: affine/non-affine, instance-norm and single-group extremes, non-square and 3D inputs for group norm; integer/anisotropic/fractional/explicit-size/downsample for upsample, with rtol=atol=0. The one gap is the float32/float64 boundary case noted in #1 — a large-extent fractional-scale config would close it.

---

Overall this is a solid, self-contained change that fixes a real partitioning cliff for diffusion UNets. My only substantive ask is confirming the float-precision behavior in #1; the rest are minor.
· branch ms/mlx-group-norm-upsample

[metascroy]

Thanks for the contribution! If CI passes, we can merge

[metascroy]

Thanks for the contribution @msluszniak!