MLX: add native_group_norm and upsample_nearest2d handlers

backends/mlx/ops.py

@@ -2756,6 +2756,246 @@ def _native_layer_norm_handler(P: MLXProgramBuilder, n: Node) -> Slot:
     return output_slots
 
 
+@REGISTRY.register(target=[torch.ops.aten.native_group_norm.default])
+def _native_group_norm_handler(P: MLXProgramBuilder, n: Node) -> Slot:
+    """Handle native_group_norm which returns (output, mean, rstd).
+
+    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
+    lets fast::layer_norm compute the normalization as a single fused kernel.
+
+    The affine parameters are applied afterwards on the original shape rather
+    than being passed to layer_norm: group norm's weight and bias are per
+    channel, while layer_norm's are per normalized element, so the two only
+    coincide when every group holds a single channel.
+
+    Only the normalized output (index 0) is computed; mean and rstd (indices 1
+    and 2) are needed only for backward.
+    """
+    unsupported = used_getitem_indices(n) & {1, 2}
+    if unsupported:
+        raise ValueError(
+            f"native_group_norm outputs {unsupported} (mean/rstd) are used, "
+            "but only the normalized output (index 0) is supported"
+        )
+
+    args = P.args(n)
+    require_args(args, 8, 8, "aten.native_group_norm")
+    require_kwargs(P.kwargs(n), set(), "aten.native_group_norm")
+    x, weight, bias, N, C, HxW, group, eps = args
+
+    for name, value in (("N", N), ("C", C), ("HxW", HxW), ("group", group)):
+        if not isinstance(value, int):
+            raise ValueError(
+                f"aten.native_group_norm requires a static {name}, got {value!r}"
+            )
+
+    x_meta = n.args[0].meta.get("val")
+    if x_meta is None:
+        raise ValueError("aten.native_group_norm requires input shape metadata")
+    if any(not isinstance(d, int) for d in x_meta.shape):
+        raise ValueError(
+            f"aten.native_group_norm requires a static input shape, "
+            f"got {tuple(x_meta.shape)}"
+        )
+    x_ndim = len(x_meta.shape)
+
+    # native_group_norm returns (output, mean, rstd) -- allocate all 3 slots
+    output_slots = P.make_or_get_slots(n)
+    out = output_slots[0]
+
+    _, flat = P.make_tmp_slot()
+    P.emit(
+        ReshapeNode(
+            x=P.slot_to_tid(x),
+            out=P.slot_to_tid(flat),
+            shape=[
+                IntOrVid.from_literal(N * group),
+                IntOrVid.from_literal((C // group) * HxW),
+            ],
+        )
+    )
+
+    _, normed = P.make_tmp_slot()
+    P.emit(
+        LayerNormNode(
+            x=P.slot_to_tid(flat),
+            out=P.slot_to_tid(normed),
+            weight=None,
+            bias=None,
+            eps=eps,
+        )
+    )
+
+    orig_shape = [IntOrVid.from_literal(int(d)) for d in x_meta.shape]
+    if weight is None and bias is None:
+        P.emit(
+            ReshapeNode(
+                x=P.slot_to_tid(normed),
+                out=P.slot_to_tid(out),
+                shape=orig_shape,
+            )
+        )
+        return output_slots
+
+    _, restored = P.make_tmp_slot()
+    P.emit(
+        ReshapeNode(
+            x=P.slot_to_tid(normed),
+            out=P.slot_to_tid(restored),
+            shape=orig_shape,
+        )
+    )
+
+    # Broadcast the per-channel affine over the batch and spatial dimensions.
+    affine_shape = [IntOrVid.from_literal(1), IntOrVid.from_literal(C)] + [
+        IntOrVid.from_literal(1)
+    ] * (x_ndim - 2)
+
+    cur = restored
+    if weight is not None:
+        _, weight_bc = P.make_tmp_slot()
+        P.emit(
+            ReshapeNode(
+                x=P.slot_to_tid(weight),
+                out=P.slot_to_tid(weight_bc),
+                shape=affine_shape,
+            )
+        )
+        if bias is None:
+            scaled = out
+        else:
+            _, scaled = P.make_tmp_slot()
+        P.emit(
+            MultiplyNode(
+                a=P.slot_to_tid(cur),
+                b=P.slot_to_tid(weight_bc),
+                out=P.slot_to_tid(scaled),
+            )
+        )
+        cur = scaled
+
+    if bias is not None:
+        _, bias_bc = P.make_tmp_slot()
+        P.emit(
+            ReshapeNode(
+                x=P.slot_to_tid(bias),
+                out=P.slot_to_tid(bias_bc),
+                shape=affine_shape,
+            )
+        )
+        P.emit(
+            AddNode(
+                a=P.slot_to_tid(cur),
+                b=P.slot_to_tid(bias_bc),
+                out=P.slot_to_tid(out),
+            )
+        )
+
+    return output_slots
+
+
+def _nearest_source_indices(in_size: int, out_size: int, scale: Optional[float]):
+    """Source index per output position for nearest-neighbour resampling.
+
+    Mirrors aten's compute_scales_value + nearest_neighbor_compute_source_index:
+    the step is ``1 / scale`` when an explicit scale factor was given and
+    ``in_size / out_size`` otherwise, and the index is truncated then clamped to
+    the last input position.
+    """
+    step = (1.0 / scale) if scale is not None else (in_size / out_size)
+    idx = (torch.arange(out_size, dtype=torch.float64) * step).to(torch.int64)
+    return torch.clamp(idx, max=in_size - 1).to(torch.int32)
+
+
+@REGISTRY.register(
+    target=[
+        torch.ops.aten.upsample_nearest2d.vec,
+        torch.ops.aten.upsample_nearest2d.default,
+    ]
+)
+def _upsample_nearest2d_handler(P: MLXProgramBuilder, n: Node) -> Slot:
+    """Nearest-neighbour 2D resampling as two gathers.
+
+    Each output position reads a source position that depends only on the static
+    input and output sizes, so both index vectors are constants and the op is
+    ``take(take(x, idx_h, -2), idx_w, -1)``. Expressing it as a gather rather
+    than a repeat also covers non-integer scale factors and downsampling.
+    """
+    args = P.args(n)
+    kwargs = P.kwargs(n)
+    x = args[0]
+
+    if ".vec" in str(n.target):
+        require_args(args, 2, 3, "aten.upsample_nearest2d.vec")
+        require_kwargs(kwargs, set(), "aten.upsample_nearest2d.vec")
+        scale_factors = args[2] if len(args) > 2 else None
+        if scale_factors is None:
+            scales = (None, None)
+        else:
+            if len(scale_factors) != 2:
+                raise ValueError(
+                    f"aten.upsample_nearest2d.vec expects 2 scale factors, "
+                    f"got {len(scale_factors)}"
+                )
+            scales = (scale_factors[0], scale_factors[1])
+    else:
+        require_args(args, 2, 4, "aten.upsample_nearest2d")
+        require_kwargs(kwargs, set(), "aten.upsample_nearest2d")
+        scales = (
+            args[2] if len(args) > 2 else None,
+            args[3] if len(args) > 3 else None,
+        )
+
+    x_meta = n.args[0].meta.get("val")
+    out_meta = n.meta.get("val")
+    if x_meta is None or out_meta is None:
+        raise ValueError("aten.upsample_nearest2d requires input and output shapes")
+
+    sizes = (x_meta.shape[-2], x_meta.shape[-1], out_meta.shape[-2], out_meta.shape[-1])
+    if any(not isinstance(d, int) for d in sizes):
+        raise ValueError(
+            f"aten.upsample_nearest2d requires static spatial sizes, got "
+            f"{tuple(x_meta.shape)} -> {tuple(out_meta.shape)}"
+        )
+    in_h, in_w, out_h, out_w = sizes
+
+    index_slots = []
+    for axis_name, in_size, out_size, scale in (
+        ("h", in_h, out_h, scales[0]),
+        ("w", in_w, out_w, scales[1]),
+    ):
+        indices = _nearest_source_indices(in_size, out_size, scale)
+        index_slots.append(
+            P.make_or_get_constant(
+                f"_upsample_nearest2d_{axis_name}_{in_size}_{out_size}_{scale}",
+                indices,
+            )
+        )
+
+    _, rows = P.make_tmp_slot()
+    P.emit(
+        TakeNode(
+            x=P.slot_to_tid(x),
+            out=P.slot_to_tid(rows),
+            index=IntOrVidOrTid.from_tid(P.slot_to_tid(index_slots[0])),
+            axis=-2,
+        )
+    )
+
+    out = P.make_or_get_slot(n)
+    P.emit(
+        TakeNode(
+            x=P.slot_to_tid(rows),
+            out=P.slot_to_tid(out),
+            index=IntOrVidOrTid.from_tid(P.slot_to_tid(index_slots[1])),
+            axis=-1,
+        )
+    )
+    return out
+
+
 @REGISTRY.register(target=[torch.ops.aten.arange.default])
 def _arange_handler(P: MLXProgramBuilder, n: Node) -> Slot:
     """Handle arange with just stop, or (start, stop) or (start, stop, step).

backends/mlx/test/test_ops.py

@@ -3332,6 +3332,136 @@ def create_inputs(self) -> Tuple[torch.Tensor, ...]:
         return (x,)
 
 
+class GroupNormModel(nn.Module):
+    """Simple model using GroupNorm."""
+
+    def __init__(
+        self,
+        num_groups: int = 8,
+        num_channels: int = 32,
+        eps: float = 1e-5,
+        affine: bool = True,
+    ):
+        super().__init__()
+        self.group_norm = nn.GroupNorm(num_groups, num_channels, eps=eps, affine=affine)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.group_norm(x)
+
+
+@register_test
+class GroupNormTest(OpTestCase):
+    """Test case for nn.GroupNorm (aten.native_group_norm)."""
+
+    name = "group_norm"
+    rtol = 1e-4
+    atol = 1e-4
+
+    def __init__(
+        self,
+        num_groups: int = 8,
+        num_channels: int = 32,
+        shape: Tuple[int, ...] = (2, 32, 8, 8),
+        eps: float = 1e-5,
+        affine: bool = True,
+        suffix: str = "",
+    ):
+        self.num_groups = num_groups
+        self.num_channels = num_channels
+        self.shape = shape
+        self.eps = eps
+        self.affine = affine
+        self.name = f"group_norm{suffix}"
+
+    @classmethod
+    def get_test_configs(cls) -> List["GroupNormTest"]:
+        return [
+            cls(),
+            # affine=False exercises the no-weight/no-bias path
+            cls(affine=False, suffix="_no_affine"),
+            # one channel per group (instance norm) and one group (all channels)
+            cls(num_groups=32, suffix="_per_channel_groups"),
+            cls(num_groups=1, suffix="_single_group"),
+            # non-square spatial extent, and a 3D (N, C, L) input
+            cls(num_groups=4, num_channels=16, shape=(1, 16, 5, 7), suffix="_odd"),
+            cls(num_groups=4, num_channels=12, shape=(2, 12, 7), suffix="_3d"),
+        ]
+
+    def create_model(self) -> nn.Module:
+        return GroupNormModel(self.num_groups, self.num_channels, self.eps, self.affine)
+
+    def create_inputs(self) -> Tuple[torch.Tensor, ...]:
+        return (torch.randn(*self.shape),)
+
+
+class UpsampleNearest2dModel(nn.Module):
+    """Nearest-neighbour resize, by scale factor or by explicit output size."""
+
+    def __init__(
+        self,
+        scale_factor: Optional[Tuple[float, float]] = None,
+        size: Optional[Tuple[int, int]] = None,
+    ):
+        super().__init__()
+        self.scale_factor = scale_factor
+        self.size = size
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return torch.nn.functional.interpolate(
+            x, size=self.size, scale_factor=self.scale_factor, mode="nearest"
+        )
+
+
+@register_test
+class UpsampleNearest2dTest(OpTestCase):
+    """Test case for aten.upsample_nearest2d."""
+
+    name = "upsample_nearest2d"
+    rtol = 0
+    atol = 0
+
+    def __init__(
+        self,
+        shape: Tuple[int, ...] = (1, 3, 4, 4),
+        scale_factor: Optional[Tuple[float, float]] = (2.0, 2.0),
+        size: Optional[Tuple[int, int]] = None,
+        suffix: str = "",
+    ):
+        self.shape = shape
+        self.scale_factor = scale_factor
+        self.size = size
+        self.name = f"upsample_nearest2d{suffix}"
+
+    @classmethod
+    def get_test_configs(cls) -> List["UpsampleNearest2dTest"]:
+        return [
+            cls(),
+            # different scale per axis
+            cls(shape=(2, 5, 3, 7), scale_factor=(3.0, 2.0), suffix="_anisotropic"),
+            # non-integer ratios, which a repeat-based lowering could not express
+            cls(shape=(1, 3, 6, 6), scale_factor=(1.5, 2.5), suffix="_fractional"),
+            cls(
+                shape=(1, 2, 5, 5),
+                scale_factor=None,
+                size=(12, 12),
+                suffix="_explicit_size",
+            ),
+            # output smaller than input
+            cls(
+                shape=(1, 4, 8, 8),
+                scale_factor=None,
+                size=(4, 4),
+                suffix="_downsample",
+            ),
+        ]
+
+    def create_model(self) -> nn.Module:
+        return UpsampleNearest2dModel(self.scale_factor, self.size)
+
+    def create_inputs(self) -> Tuple[torch.Tensor, ...]:
+        return (torch.randn(*self.shape),)
+
+
 class Conv1dModel(nn.Module):
     """Simple model using Conv1d."""