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

[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.

❌ 5 Pending, 1 Unclassified Failure

As of commit bb67233 with merge base fbd4bbf ():

UNCLASSIFIED FAILURE - DrCI could not classify the following job because the workflow did not run on the merge base. The failure may be pre-existing on trunk or introduced by this PR:

• MLX / test-mlx-parakeet / test-mlx-parakeet (gh) (this job did not run on the merge base, so DrCI cannot tell whether the failure is pre-existing)
  RuntimeError: Command bash /Users/runner/work/_temp/exec_script failed with exit code 1

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.