[Relax][Frontend][ONNX] Add GroupNormalization support

[napronald]

Summary

Adds ONNX frontend support for GroupNormalization by mapping it to the existing relax.op.nn.group_norm.

Supports opset 18 per-group scale/bias expansion, opset 21 per-channel scale/bias, and stash_type cast behavior.

Testing

Includes structural checks for opset 18, opset 21, rank-3 inputs, and fp16 stash_type paths.

[gemini-code-assist]

Code Review

This pull request adds support for converting ONNX GroupNormalization nodes (opsets 18 and 21) into equivalent Relax expressions, along with comprehensive unit tests to verify correctness. The feedback suggests optimizing the opset 18 implementation by resolving scale and bias to constants at import time using get_constant. If they are constants, they can be expanded directly using NumPy (_np.repeat), which avoids adding redundant reshape and broadcast_to operators to the generated Relax graph.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[gemini-code-assist]

We can use get_constant to resolve scale and bias to constants if they are initializers. This allows us to perform the per-group to per-channel expansion at import time using NumPy, avoiding redundant reshape and broadcast_to operators in the Relax graph.

Suggested change

	data = inputs[0]
	scale = inputs[1]
	bias = inputs[2]
	data = inputs[0]
	scale = get_constant(inputs[1], params)
	bias = get_constant(inputs[2], params)

[gemini-code-assist]

If scale and bias are resolved to constants, we can expand them directly using _np.repeat at import time. This simplifies the generated Relax graph by eliminating unnecessary reshape and broadcast_to operations.

Suggested change

	scale = relax.op.reshape(scale, [num_groups, 1])
	scale = relax.op.broadcast_to(scale, [num_groups, channels_per_group])
	scale = relax.op.reshape(scale, [channels])
	bias = relax.op.reshape(bias, [num_groups, 1])
	bias = relax.op.broadcast_to(bias, [num_groups, channels_per_group])
	bias = relax.op.reshape(bias, [channels])
	if isinstance(scale, relax.Constant):
	scale = relax.const(_np.repeat(scale.data.numpy(), channels_per_group), scale.ty.dtype)
	else:
	scale = relax.op.reshape(scale, [num_groups, 1])
	scale = relax.op.broadcast_to(scale, [num_groups, channels_per_group])
	scale = relax.op.reshape(scale, [channels])
	if isinstance(bias, relax.Constant):
	bias = relax.const(_np.repeat(bias.data.numpy(), channels_per_group), bias.ty.dtype)
	else:
	bias = relax.op.reshape(bias, [num_groups, 1])
	bias = relax.op.broadcast_to(bias, [num_groups, channels_per_group])
	bias = relax.op.reshape(bias, [channels])

[tlopex]

Thanks for the PR. There are still some issues should be resolved

1. For opset 21, stash_type only controls the precision of the normalization stage. The normalized result must be cast back to the input dtype before applying scale and bias. The current implementation casts data, scale, and bias to float32, performs the entire group norm including the affine stage in float32, and only casts the final result back. For fp16 with the default stash_type=1, I reproduced 4/16 outputs differing from onnx.reference.ReferenceEvaluator, with a maximum absolute difference of 0.0625.

2. stash_type is a TensorProto.DataType, not a boolean flag. 0 is UNDEFINED, while valid floating-point values include FLOAT=1, FLOAT16=10, DOUBLE=11, and BFLOAT16=16. The current [0, 1] check accepts an invalid value and rejects valid values. If only FLOAT is supported for now, it would be better to reject the other values explicitly.

3. Opset 18 fp16 also has a precision mismatch. ONNX-18 performs the computation in the input dtype, while TVM's relax.nn.group_norm internally promotes fp16 normalization to float32. I reproduced 12/16 outputs differing from the ONNX reference implementation.

Could we fix them and update the tests accordingly? Thanks