MLX: add native_group_norm and upsample_nearest2d handlers by msluszniak · Pull Request #22050 · pytorch/executorch

msluszniak · 2026-08-22T17:10:38Z

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 · 2026-08-22T17:10:41Z

🔗 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 · 2026-08-22T17:10:45Z

The committers listed above are authorized under a signed CLA.

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

github-actions · 2026-08-22T17:11:27Z

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.

The MLX backend had 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 instead of merely slowing it down: the SDXS-512-DreamShaper UNet partitions into 28 delegate subgraphs, leaving 25 native_group_norm and 2 upsample_nearest2d nodes on the CPU, and each boundary crossing leaves and re-enters the MLX runtime. Both ops lower to primitives the backend already has, so this needs no schema or runtime change. native_group_norm normalizes each group of C / group channels together with all of their spatial positions, so reshaping the input to (N * group, (C / group) * HxW) puts exactly that set on the last axis and fast::layer_norm computes it as one fused kernel. The affine parameters are applied afterwards on the original shape rather than being handed 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. 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), 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. With the handlers registered the same UNet lowers to a single delegate subgraph with nothing left on the CPU. Adds op tests for both: group norm over affine and non-affine, one channel per group and one group for all channels, non-square spatial extents and a 3D (N, C, L) input; upsampling over integer, anisotropic and fractional scale factors, explicit output sizes and downsampling. All 11 configurations match eager through the MLX runtime, the upsample ones bit-exactly. Fixespytorch#22017

metascroy · 2026-08-25T16:27:49Z

@claude review this code

claude · 2026-08-25T16:28:48Z

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= ... # floatidx= (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"instr(n.target):

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

ifn.targetistorch.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 · 2026-08-26T00:05:00Z

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

metascroy · 2026-08-26T19:33:51Z

Thanks for the contribution @msluszniak!

msluszniak · 2026-08-22T17:10:38Z

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 · 2026-08-22T17:10:41Z

🔗 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 · 2026-08-22T17:10:45Z

The committers listed above are authorized under a signed CLA.

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

github-actions · 2026-08-22T17:11:27Z

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.

The MLX backend had 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 instead of merely slowing it down: the SDXS-512-DreamShaper UNet partitions into 28 delegate subgraphs, leaving 25 native_group_norm and 2 upsample_nearest2d nodes on the CPU, and each boundary crossing leaves and re-enters the MLX runtime. Both ops lower to primitives the backend already has, so this needs no schema or runtime change. native_group_norm normalizes each group of C / group channels together with all of their spatial positions, so reshaping the input to (N * group, (C / group) * HxW) puts exactly that set on the last axis and fast::layer_norm computes it as one fused kernel. The affine parameters are applied afterwards on the original shape rather than being handed 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. 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), 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. With the handlers registered the same UNet lowers to a single delegate subgraph with nothing left on the CPU. Adds op tests for both: group norm over affine and non-affine, one channel per group and one group for all channels, non-square spatial extents and a 3D (N, C, L) input; upsampling over integer, anisotropic and fractional scale factors, explicit output sizes and downsampling. All 11 configurations match eager through the MLX runtime, the upsample ones bit-exactly. Fixespytorch#22017

metascroy · 2026-08-25T16:27:49Z

@claude review this code

claude · 2026-08-25T16:28:48Z

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= ... # floatidx= (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"instr(n.target):

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

ifn.targetistorch.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 · 2026-08-26T00:05:00Z

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

metascroy · 2026-08-26T19:33:51Z

Thanks for the contribution @msluszniak!

msluszniak · 2026-08-22T17:10:38Z

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 · 2026-08-22T17:10:41Z

🔗 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 · 2026-08-22T17:10:45Z

The committers listed above are authorized under a signed CLA.

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

github-actions · 2026-08-22T17:11:27Z

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.

The MLX backend had 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 instead of merely slowing it down: the SDXS-512-DreamShaper UNet partitions into 28 delegate subgraphs, leaving 25 native_group_norm and 2 upsample_nearest2d nodes on the CPU, and each boundary crossing leaves and re-enters the MLX runtime. Both ops lower to primitives the backend already has, so this needs no schema or runtime change. native_group_norm normalizes each group of C / group channels together with all of their spatial positions, so reshaping the input to (N * group, (C / group) * HxW) puts exactly that set on the last axis and fast::layer_norm computes it as one fused kernel. The affine parameters are applied afterwards on the original shape rather than being handed 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. 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), 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. With the handlers registered the same UNet lowers to a single delegate subgraph with nothing left on the CPU. Adds op tests for both: group norm over affine and non-affine, one channel per group and one group for all channels, non-square spatial extents and a 3D (N, C, L) input; upsampling over integer, anisotropic and fractional scale factors, explicit output sizes and downsampling. All 11 configurations match eager through the MLX runtime, the upsample ones bit-exactly. Fixespytorch#22017

metascroy · 2026-08-25T16:27:49Z

@claude review this code

claude · 2026-08-25T16:28:48Z

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= ... # floatidx= (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"instr(n.target):

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

ifn.targetistorch.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 · 2026-08-26T00:05:00Z

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

metascroy · 2026-08-26T19:33:51Z

Thanks for the contribution @msluszniak!

msluszniak · 2026-08-22T17:10:38Z

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 · 2026-08-22T17:10:41Z

🔗 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 · 2026-08-22T17:10:45Z

The committers listed above are authorized under a signed CLA.

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

github-actions · 2026-08-22T17:11:27Z

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.

The MLX backend had 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 instead of merely slowing it down: the SDXS-512-DreamShaper UNet partitions into 28 delegate subgraphs, leaving 25 native_group_norm and 2 upsample_nearest2d nodes on the CPU, and each boundary crossing leaves and re-enters the MLX runtime. Both ops lower to primitives the backend already has, so this needs no schema or runtime change. native_group_norm normalizes each group of C / group channels together with all of their spatial positions, so reshaping the input to (N * group, (C / group) * HxW) puts exactly that set on the last axis and fast::layer_norm computes it as one fused kernel. The affine parameters are applied afterwards on the original shape rather than being handed 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. 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), 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. With the handlers registered the same UNet lowers to a single delegate subgraph with nothing left on the CPU. Adds op tests for both: group norm over affine and non-affine, one channel per group and one group for all channels, non-square spatial extents and a 3D (N, C, L) input; upsampling over integer, anisotropic and fractional scale factors, explicit output sizes and downsampling. All 11 configurations match eager through the MLX runtime, the upsample ones bit-exactly. Fixespytorch#22017

metascroy · 2026-08-25T16:27:49Z

@claude review this code

claude · 2026-08-25T16:28:48Z

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= ... # floatidx= (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"instr(n.target):

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

ifn.targetistorch.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 · 2026-08-26T00:05:00Z

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

metascroy · 2026-08-26T19:33:51Z

Thanks for the contribution @msluszniak!

msluszniak · 2026-08-22T17:10:38Z

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 · 2026-08-22T17:10:41Z

🔗 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 · 2026-08-22T17:10:45Z

The committers listed above are authorized under a signed CLA.

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

github-actions · 2026-08-22T17:11:27Z

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.

The MLX backend had 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 instead of merely slowing it down: the SDXS-512-DreamShaper UNet partitions into 28 delegate subgraphs, leaving 25 native_group_norm and 2 upsample_nearest2d nodes on the CPU, and each boundary crossing leaves and re-enters the MLX runtime. Both ops lower to primitives the backend already has, so this needs no schema or runtime change. native_group_norm normalizes each group of C / group channels together with all of their spatial positions, so reshaping the input to (N * group, (C / group) * HxW) puts exactly that set on the last axis and fast::layer_norm computes it as one fused kernel. The affine parameters are applied afterwards on the original shape rather than being handed 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. 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), 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. With the handlers registered the same UNet lowers to a single delegate subgraph with nothing left on the CPU. Adds op tests for both: group norm over affine and non-affine, one channel per group and one group for all channels, non-square spatial extents and a 3D (N, C, L) input; upsampling over integer, anisotropic and fractional scale factors, explicit output sizes and downsampling. All 11 configurations match eager through the MLX runtime, the upsample ones bit-exactly. Fixespytorch#22017

metascroy · 2026-08-25T16:27:49Z

@claude review this code

claude · 2026-08-25T16:28:48Z

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= ... # floatidx= (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"instr(n.target):

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

ifn.targetistorch.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 · 2026-08-26T00:05:00Z

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

metascroy · 2026-08-26T19:33:51Z

Thanks for the contribution @msluszniak!

msluszniak · 2026-08-22T17:10:38Z

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 · 2026-08-22T17:10:41Z

🔗 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 · 2026-08-22T17:10:45Z

The committers listed above are authorized under a signed CLA.

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

github-actions · 2026-08-22T17:11:27Z

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.

The MLX backend had 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 instead of merely slowing it down: the SDXS-512-DreamShaper UNet partitions into 28 delegate subgraphs, leaving 25 native_group_norm and 2 upsample_nearest2d nodes on the CPU, and each boundary crossing leaves and re-enters the MLX runtime. Both ops lower to primitives the backend already has, so this needs no schema or runtime change. native_group_norm normalizes each group of C / group channels together with all of their spatial positions, so reshaping the input to (N * group, (C / group) * HxW) puts exactly that set on the last axis and fast::layer_norm computes it as one fused kernel. The affine parameters are applied afterwards on the original shape rather than being handed 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. 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), 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. With the handlers registered the same UNet lowers to a single delegate subgraph with nothing left on the CPU. Adds op tests for both: group norm over affine and non-affine, one channel per group and one group for all channels, non-square spatial extents and a 3D (N, C, L) input; upsampling over integer, anisotropic and fractional scale factors, explicit output sizes and downsampling. All 11 configurations match eager through the MLX runtime, the upsample ones bit-exactly. Fixespytorch#22017

metascroy · 2026-08-25T16:27:49Z

@claude review this code

claude · 2026-08-25T16:28:48Z

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= ... # floatidx= (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"instr(n.target):

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

ifn.targetistorch.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 · 2026-08-26T00:05:00Z

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

metascroy · 2026-08-26T19:33:51Z

Thanks for the contribution @msluszniak!

msluszniak · 2026-08-22T17:10:38Z

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 · 2026-08-22T17:10:41Z

🔗 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 · 2026-08-22T17:10:45Z

The committers listed above are authorized under a signed CLA.

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

github-actions · 2026-08-22T17:11:27Z

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.

The MLX backend had 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 instead of merely slowing it down: the SDXS-512-DreamShaper UNet partitions into 28 delegate subgraphs, leaving 25 native_group_norm and 2 upsample_nearest2d nodes on the CPU, and each boundary crossing leaves and re-enters the MLX runtime. Both ops lower to primitives the backend already has, so this needs no schema or runtime change. native_group_norm normalizes each group of C / group channels together with all of their spatial positions, so reshaping the input to (N * group, (C / group) * HxW) puts exactly that set on the last axis and fast::layer_norm computes it as one fused kernel. The affine parameters are applied afterwards on the original shape rather than being handed 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. 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), 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. With the handlers registered the same UNet lowers to a single delegate subgraph with nothing left on the CPU. Adds op tests for both: group norm over affine and non-affine, one channel per group and one group for all channels, non-square spatial extents and a 3D (N, C, L) input; upsampling over integer, anisotropic and fractional scale factors, explicit output sizes and downsampling. All 11 configurations match eager through the MLX runtime, the upsample ones bit-exactly. Fixespytorch#22017

metascroy · 2026-08-25T16:27:49Z

@claude review this code

claude · 2026-08-25T16:28:48Z

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= ... # floatidx= (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"instr(n.target):

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

ifn.targetistorch.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 · 2026-08-26T00:05:00Z

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

metascroy · 2026-08-26T19:33:51Z

Thanks for the contribution @msluszniak!

msluszniak · 2026-08-22T17:10:38Z

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 · 2026-08-22T17:10:41Z

🔗 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 · 2026-08-22T17:10:45Z

The committers listed above are authorized under a signed CLA.

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

github-actions · 2026-08-22T17:11:27Z

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.

The MLX backend had 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 instead of merely slowing it down: the SDXS-512-DreamShaper UNet partitions into 28 delegate subgraphs, leaving 25 native_group_norm and 2 upsample_nearest2d nodes on the CPU, and each boundary crossing leaves and re-enters the MLX runtime. Both ops lower to primitives the backend already has, so this needs no schema or runtime change. native_group_norm normalizes each group of C / group channels together with all of their spatial positions, so reshaping the input to (N * group, (C / group) * HxW) puts exactly that set on the last axis and fast::layer_norm computes it as one fused kernel. The affine parameters are applied afterwards on the original shape rather than being handed 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. 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), 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. With the handlers registered the same UNet lowers to a single delegate subgraph with nothing left on the CPU. Adds op tests for both: group norm over affine and non-affine, one channel per group and one group for all channels, non-square spatial extents and a 3D (N, C, L) input; upsampling over integer, anisotropic and fractional scale factors, explicit output sizes and downsampling. All 11 configurations match eager through the MLX runtime, the upsample ones bit-exactly. Fixespytorch#22017

metascroy · 2026-08-25T16:27:49Z

@claude review this code

claude · 2026-08-25T16:28:48Z

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= ... # floatidx= (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"instr(n.target):

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

ifn.targetistorch.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 · 2026-08-26T00:05:00Z

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

metascroy · 2026-08-26T19:33:51Z

Thanks for the contribution @msluszniak!

meta-claBot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label Aug 22, 2026

msluszniak force-pushed the ms/mlx-group-norm-upsample branch from 9d8cb4a to bb67233Compare August 22, 2026 17:15

JacobSzwejbka requested a review from metascroy August 24, 2026 17:26

metascroy added the ciflow/mlx label Aug 26, 2026

metascroy approved these changes Aug 26, 2026
View reviewed changes

metascroy merged commit 0cf9774 into pytorch:mainAug 26, 2026
259 of 260 checks passed

nil-is-all added the module: mlx Issues related to MLX Backend: Metal-accelerated inference on Apple Silicon label Aug 31, 2026

nil-is-all assigned metascroy Aug 31, 2026

meta-claBot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label Aug 22, 2026

msluszniak force-pushed the ms/mlx-group-norm-upsample branch from 9d8cb4a to bb67233Compare August 22, 2026 17:15

JacobSzwejbka requested a review from metascroy August 24, 2026 17:26

metascroy added the ciflow/mlx label Aug 26, 2026

metascroy approved these changes Aug 26, 2026
View reviewed changes

metascroy merged commit 0cf9774 into pytorch:mainAug 26, 2026
259 of 260 checks passed

nil-is-all added the module: mlx Issues related to MLX Backend: Metal-accelerated inference on Apple Silicon label Aug 31, 2026

nil-is-all assigned metascroy Aug 31, 2026

meta-claBot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label Aug 22, 2026

msluszniak force-pushed the ms/mlx-group-norm-upsample branch from 9d8cb4a to bb67233Compare August 22, 2026 17:15

JacobSzwejbka requested a review from metascroy August 24, 2026 17:26

metascroy added the ciflow/mlx label Aug 26, 2026

metascroy approved these changes Aug 26, 2026
View reviewed changes

metascroy merged commit 0cf9774 into pytorch:mainAug 26, 2026
259 of 260 checks passed

nil-is-all added the module: mlx Issues related to MLX Backend: Metal-accelerated inference on Apple Silicon label Aug 31, 2026

nil-is-all assigned metascroy Aug 31, 2026

meta-claBot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label Aug 22, 2026

msluszniak force-pushed the ms/mlx-group-norm-upsample branch from 9d8cb4a to bb67233Compare August 22, 2026 17:15

JacobSzwejbka requested a review from metascroy August 24, 2026 17:26

metascroy added the ciflow/mlx label Aug 26, 2026

metascroy approved these changes Aug 26, 2026
View reviewed changes

metascroy merged commit 0cf9774 into pytorch:mainAug 26, 2026
259 of 260 checks passed

nil-is-all added the module: mlx Issues related to MLX Backend: Metal-accelerated inference on Apple Silicon label Aug 31, 2026

nil-is-all assigned metascroy Aug 31, 2026

meta-claBot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label Aug 22, 2026

msluszniak force-pushed the ms/mlx-group-norm-upsample branch from 9d8cb4a to bb67233Compare August 22, 2026 17:15

JacobSzwejbka requested a review from metascroy August 24, 2026 17:26

metascroy added the ciflow/mlx label Aug 26, 2026

metascroy approved these changes Aug 26, 2026
View reviewed changes

metascroy merged commit 0cf9774 into pytorch:mainAug 26, 2026
259 of 260 checks passed

nil-is-all added the module: mlx Issues related to MLX Backend: Metal-accelerated inference on Apple Silicon label Aug 31, 2026

nil-is-all assigned metascroy Aug 31, 2026

meta-claBot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label Aug 22, 2026

msluszniak force-pushed the ms/mlx-group-norm-upsample branch from 9d8cb4a to bb67233Compare August 22, 2026 17:15

JacobSzwejbka requested a review from metascroy August 24, 2026 17:26

metascroy added the ciflow/mlx label Aug 26, 2026

metascroy approved these changes Aug 26, 2026
View reviewed changes

metascroy merged commit 0cf9774 into pytorch:mainAug 26, 2026
259 of 260 checks passed

nil-is-all added the module: mlx Issues related to MLX Backend: Metal-accelerated inference on Apple Silicon label Aug 31, 2026

nil-is-all assigned metascroy Aug 31, 2026

meta-claBot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label Aug 22, 2026

msluszniak force-pushed the ms/mlx-group-norm-upsample branch from 9d8cb4a to bb67233Compare August 22, 2026 17:15

JacobSzwejbka requested a review from metascroy August 24, 2026 17:26

metascroy added the ciflow/mlx label Aug 26, 2026

metascroy approved these changes Aug 26, 2026
View reviewed changes

metascroy merged commit 0cf9774 into pytorch:mainAug 26, 2026
259 of 260 checks passed

nil-is-all added the module: mlx Issues related to MLX Backend: Metal-accelerated inference on Apple Silicon label Aug 31, 2026

nil-is-all assigned metascroy Aug 31, 2026

meta-claBot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label Aug 22, 2026

msluszniak force-pushed the ms/mlx-group-norm-upsample branch from 9d8cb4a to bb67233Compare August 22, 2026 17:15

JacobSzwejbka requested a review from metascroy August 24, 2026 17:26

metascroy added the ciflow/mlx label Aug 26, 2026

metascroy approved these changes Aug 26, 2026
View reviewed changes

metascroy merged commit 0cf9774 into pytorch:mainAug 26, 2026
259 of 260 checks passed

nil-is-all added the module: mlx Issues related to MLX Backend: Metal-accelerated inference on Apple Silicon label Aug 31, 2026

nil-is-all assigned metascroy Aug 31, 2026

Uh oh!

Conversation

msluszniak commented Aug 22, 2026• edited by pytorch-bot BotLoading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Approach

Test plan

Uh oh!

pytorch-botBot commented Aug 22, 2026• edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

🔗 Helpful Links

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

⚠️ 1 Awaiting Approval

Uh oh!

linux-foundation-easyclaBot commented Aug 22, 2026• edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

github-actionsBot commented Aug 22, 2026

This PR needs a release notes: label

Uh oh!

metascroy commented Aug 25, 2026

Uh oh!

claudeBot commented Aug 25, 2026• edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Code Review

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

2. native_group_norm reshape assumes NCHW-contiguous input

3. .vec detection via string match — minor robustness

Things I checked that look correct

Test coverage

Uh oh!

metascroy commented Aug 26, 2026

Uh oh!

Uh oh!

metascroy commented Aug 26, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Uh oh!

Conversation

msluszniak commented Aug 22, 2026• edited by pytorch-bot BotLoading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Approach

Test plan

Uh oh!

pytorch-botBot commented Aug 22, 2026• edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

🔗 Helpful Links

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

⚠️ 1 Awaiting Approval

Uh oh!

linux-foundation-easyclaBot commented Aug 22, 2026• edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

github-actionsBot commented Aug 22, 2026

This PR needs a release notes: label

Uh oh!

metascroy commented Aug 25, 2026

Uh oh!

claudeBot commented Aug 25, 2026• edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Code Review

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

2. native_group_norm reshape assumes NCHW-contiguous input

3. .vec detection via string match — minor robustness

Things I checked that look correct

Test coverage

Uh oh!

metascroy commented Aug 26, 2026

Uh oh!

Uh oh!

metascroy commented Aug 26, 2026

Uh oh!

msluszniak commented Aug 22, 2026•
edited by pytorch-bot Bot
Loading

pytorch-botBot commented Aug 22, 2026•
edited
Loading

linux-foundation-easyclaBot commented Aug 22, 2026•
edited
Loading

This PR needs a `release notes:` label

claudeBot commented Aug 25, 2026•
edited
Loading

1. Index computation precision vs. aten (potential correctness edge) — `ops.py_nearest_source_indices`

2. `native_group_norm` reshape assumes NCHW-contiguous input

3. `.vec` detection via string match — minor robustness

msluszniak commented Aug 22, 2026•
edited by pytorch-bot Bot
Loading

pytorch-botBot commented Aug 22, 2026•
edited
Loading

linux-foundation-easyclaBot commented Aug 22, 2026•
edited
Loading

This PR needs a `release notes:` label

claudeBot commented Aug 25, 2026•
edited
Loading

1. Index computation precision vs. aten (potential correctness edge) — `ops.py_nearest_source_indices`

2. `native_group_norm` reshape assumes NCHW-contiguous input

3. `.vec` detection via string match — minor robustness

msluszniak commented Aug 22, 2026•
edited by pytorch-bot Bot
Loading

pytorch-botBot commented Aug 22, 2026•
edited
Loading

linux-foundation-easyclaBot commented Aug 22, 2026•
edited
Loading

This PR needs a `release notes:` label

claudeBot commented Aug 25, 2026•
edited
Loading

1. Index computation precision vs. aten (potential correctness edge) — `ops.py_nearest_source_indices`

2. `native_group_norm` reshape assumes NCHW-contiguous input

3. `.vec` detection via string match — minor robustness