DecomposeGroupNormPass(ArmPass) fails

Trying to run prepare_pt2e on a model with torch.nn.GroupNorm to the Ethos-u backend fails in DecomposeGroupNormPass(ArmPass).

Code example:

class GNModule(torch.nn.Module):
    def __init__(self, num_channels: int, num_groups: int):
        super(GNModule, self).__init__()
        self.gn = torch.nn.GroupNorm(num_groups=num_groups, num_channels=num_channels, affine=False)

    def forward(self, x):
        return self.gn(x)


def quantize_model(model, batches, backend="ethos-u", backend_cfg=None) -> torch.fx.GraphModule:
    if backend == "ethos-u":
        compile_spec = EthosUCompileSpec(**backend_cfg.get("compile_spec", {}))
        quantizer = EthosUQuantizer(compile_spec)
        operator_config = get_arm_symmetric_qconfig(is_per_channel=backend_cfg.get("is_per_channel", True))
    elif backend == "xnnpack":
        compile_spec = None
        quantizer = XNNPACKQuantizer()
        operator_config = get_xnnpack_symmetric_qconfig(is_per_channel=backend_cfg.get("is_per_channel", True))
    else:
        raise ValueError(f"Unsupported backend: {backend}")

    dynamic_shapes = {"x": {0: Dim("batch", min=1)}}

    # Generate dummy input for export with batch size > 1 so the export captures dynamic batch size
    dummy_input = torch.randn((2,)+batches[0].shape[1:]).to(batches[0].device)
    exported_program = torch.export.export(model, (dummy_input,), dynamic_shapes=dynamic_shapes)
    graph_module = exported_program.module(check_guards=False)

    quantizer.set_global(operator_config)
    prepared = prepare_pt2e(graph_module, quantizer)

    with torch.no_grad():
        for data_batch in tqdm(batches, desc="Quantization calibration"):
            prepared(data_batch)

    # print_pt2e_quant_details(prepared)
    prepared = prepared.to("cpu")

    return prepared, compile_spec


model = GNModule(num_channels=6, num_groups=3)
input_tensor = torch.randn(20, 6, 10, 10)  # batch_size=20, channels=6, height=10, width=10
output = model(input_tensor)
backend = "ethos-u"
# backend = "xnnpack"
compile_spec = {'target': 'ethos-u55-128',
                'system_config': 'Ethos_U55_Deep_Embedded',
                'memory_mode': 'Shared_Sram',
                'extra_flags': ["--output-format=raw", "--debug-force-regor", "--verbose-all"]}
qmodel, _ = quantize_model(model, [input_tensor], backend=backend, backend_cfg={
    'compile_spec': compile_spec, 'is_per_channel': True})

setting backend to xnnpack doesn't fail.

Environtment:

PyTorch version: 2.9.0
Is debug build: False
CUDA used to build PyTorch: None
ROCM used to build PyTorch: N/A

OS: macOS 15.5 (arm64)
GCC version: Could not collect
Clang version: 17.0.0 (clang-1700.0.13.5)
CMake version: Could not collect
Libc version: N/A

Python version: 3.12.10 (v3.12.10:0cc81280367, Apr 8 2025, 08:46:59) [Clang 13.0.0 (clang-1300.0.29.30)] (64-bit runtime)
Python platform: macOS-15.5-arm64-arm-64bit
Is CUDA available: False
CUDA runtime version: No CUDA
CUDA_MODULE_LOADING set to: N/A
GPU models and configuration: No CUDA
Nvidia driver version: No CUDA
cuDNN version: No CUDA
Is XPU available: False
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

CPU:
Apple M4 Max

Versions of relevant libraries:
[pip3] executorch==1.0.0
[pip3] numpy==2.3.4
[pip3] torch==2.9.0
[pip3] torchao==0.14.0
[pip3] torchaudio==2.9.0
[pip3] torchcodec==0.8.1
[pip3] torchvision==0.24.0

cc @freddan80 @per @zingo @oscarandersson8218 @digantdesai

Run the supplied GroupNorm and prepare_pt2e reproduction with the Ethos-U backend, then inspect DecomposeGroupNormPass(ArmPass), the named failing entry point. Done means the example prepares successfully for Ethos-U with torch.nn.GroupNorm, while the existing XNNPACK path remains unaffected.