This is an automated email from the ASF dual-hosted git repository.
tlopex pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/main by this push:
new 6b7380e6e2 [Relax][Frontend][ONNX] Add GroupNormalization support
(#19907)
6b7380e6e2 is described below
commit 6b7380e6e2d01f295fc330242163f47264ffcc06
Author: Ronald Nap <[email protected]>
AuthorDate: Sat Jul 11 21:24:34 2026 -0700
[Relax][Frontend][ONNX] Add GroupNormalization support (#19907)
## 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.
---
python/tvm/relax/frontend/onnx/onnx_frontend.py | 134 +++++++++++++
tests/python/relax/test_frontend_onnx.py | 248 ++++++++++++++++++++++++
2 files changed, 382 insertions(+)
diff --git a/python/tvm/relax/frontend/onnx/onnx_frontend.py
b/python/tvm/relax/frontend/onnx/onnx_frontend.py
index da65cc7bee..d340c696b4 100644
--- a/python/tvm/relax/frontend/onnx/onnx_frontend.py
+++ b/python/tvm/relax/frontend/onnx/onnx_frontend.py
@@ -4005,6 +4005,139 @@ class RMSNormalization(OnnxOpConverter):
return output
+class GroupNormalization(OnnxOpConverter):
+ """Converts an onnx GroupNormalization node into an equivalent Relax
expression"""
+
+ @classmethod
+ def _impl_v18(cls, bb, inputs, attr, params):
+ data = inputs[0]
+ scale = inputs[1]
+ bias = inputs[2]
+ num_groups = attr["num_groups"]
+ epsilon = attr.get("epsilon", 1e-05)
+
+ ndim = _get_known_tensor_rank(data)
+ if ndim is None:
+ raise ValueError("GroupNormalization requires a statically known
input rank.")
+
+ ty = data.ty
+ if not isinstance(ty, relax.TensorType) or len(ty.shape) < 2:
+ raise ValueError(
+ "GroupNormalization-18 requires a statically typed input with
rank >= 2."
+ )
+
+ input_dtype = ty.dtype
+ if input_dtype != "float32":
+ raise ValueError("GroupNormalization-18 currently only supports
float32 inputs.")
+
+ if num_groups <= 0:
+ raise ValueError(
+ f"GroupNormalization requires num_groups to be positive, got
{num_groups}."
+ )
+
+ channel_dim = ty.shape[1]
+ if not isinstance(channel_dim, tirx.IntImm):
+ raise ValueError(
+ "GroupNormalization-18 requires a statically known channel
count "
+ "to expand per-group scale/bias to per-channel."
+ )
+
+ channels = int(channel_dim)
+ if channels % num_groups != 0:
+ raise ValueError(
+ f"GroupNormalization requires num_groups to divide channel
count, "
+ f"but got C={channels} and num_groups={num_groups}."
+ )
+
+ channels_per_group = channels // num_groups
+
+ 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])
+
+ axes = list(range(2, ndim))
+ return relax.op.nn.group_norm(
+ data, scale, bias, num_groups, channel_axis=1, axes=axes,
epsilon=epsilon
+ )
+
+ @classmethod
+ def _impl_v21(cls, bb, inputs, attr, params):
+ data = inputs[0]
+ scale = inputs[1]
+ bias = inputs[2]
+ num_groups = attr["num_groups"]
+ epsilon = attr.get("epsilon", 1e-05)
+ stash_type = attr.get("stash_type", 1)
+
+ if stash_type != 1:
+ raise ValueError(
+ f"GroupNormalization currently only supports stash_type=1
(FLOAT), "
+ f"but got stash_type={stash_type}."
+ )
+
+ ndim = _get_known_tensor_rank(data)
+ if ndim is None:
+ raise ValueError("GroupNormalization requires a statically known
input rank.")
+
+ ty = data.ty
+ if not isinstance(ty, relax.TensorType) or len(ty.shape) < 2:
+ raise ValueError("GroupNormalization requires a statically typed
input with rank >= 2.")
+
+ if num_groups <= 0:
+ raise ValueError(
+ f"GroupNormalization requires num_groups to be positive, got
{num_groups}."
+ )
+
+ channel_dim = ty.shape[1]
+ if isinstance(channel_dim, tirx.IntImm):
+ channels = int(channel_dim)
+ if channels % num_groups != 0:
+ raise ValueError(
+ f"GroupNormalization requires num_groups to divide channel
count, "
+ f"but got C={channels} and num_groups={num_groups}."
+ )
+
+ axes = list(range(2, ndim))
+ input_dtype = ty.dtype
+
+ orig_scale = scale
+ orig_bias = bias
+
+ if input_dtype != "float32":
+ data = relax.op.astype(data, "float32")
+ scale = relax.op.astype(scale, "float32")
+ bias = relax.op.astype(bias, "float32")
+
+ norm_scale = relax.op.ones_like(scale)
+ norm_bias = relax.op.zeros_like(bias)
+
+ output = relax.op.nn.group_norm(
+ data,
+ norm_scale,
+ norm_bias,
+ num_groups,
+ channel_axis=1,
+ axes=axes,
+ epsilon=epsilon,
+ center=False,
+ scale=False,
+ )
+
+ if input_dtype != "float32":
+ output = relax.op.astype(output, input_dtype)
+
+ affine_shape = [channel_dim] + [1] * (ndim - 2)
+ orig_scale = relax.op.reshape(orig_scale, affine_shape)
+ orig_bias = relax.op.reshape(orig_bias, affine_shape)
+ output = relax.op.multiply(output, orig_scale)
+ output = relax.op.add(output, orig_bias)
+ return output
+
+
class ReduceMax(OnnxOpConverter):
"""Converts an onnx ReduceMax node into an equivalent Relax expression."""
@@ -5385,6 +5518,7 @@ def _get_convert_map():
"BatchNormalization": BatchNormalization,
"LayerNormalization": LayerNormalization,
"RMSNormalization": RMSNormalization,
+ "GroupNormalization": GroupNormalization,
"SkipLayerNormalization": SkipLayerNormalization,
"EmbedLayerNormalization": EmbedLayerNormalization,
"InstanceNormalization": InstanceNormalization,
diff --git a/tests/python/relax/test_frontend_onnx.py
b/tests/python/relax/test_frontend_onnx.py
index 35cf699dba..126c909983 100644
--- a/tests/python/relax/test_frontend_onnx.py
+++ b/tests/python/relax/test_frontend_onnx.py
@@ -4367,6 +4367,254 @@ def test_rms_norm():
check_correctness(model, opset=23, rtol=1e-2, atol=1e-2)
+def _make_group_norm_expected_ir(
+ input_shape: list[int],
+ scale_shape: list[int],
+ bias_shape: list[int],
+ num_groups: int,
+ opset: int = 21,
+ dtype: str = "float32",
+ stash_type: int = 1,
+):
+ input_shape = tuple(input_shape)
+ scale_shape = tuple(scale_shape)
+ bias_shape = tuple(bias_shape)
+ axes = list(range(2, len(input_shape)))
+ epsilon = float(np.float32(1e-5))
+ affine_shape = (input_shape[1],) + (1,) * (len(input_shape) - 2)
+
+ if opset == 18:
+ channels = input_shape[1]
+ channels_per_group = channels // num_groups
+
+ @I.ir_module
+ class ExpectedGroupNormOpset18:
+ @R.function
+ def main(
+ input: R.Tensor(input_shape, dtype=dtype),
+ scale: R.Tensor(scale_shape, dtype=dtype),
+ bias: R.Tensor(bias_shape, dtype=dtype),
+ ) -> R.Tensor(input_shape, dtype=dtype):
+ R.func_attr({"num_input": 3})
+ with R.dataflow():
+ lv: R.Tensor((num_groups, 1), dtype=dtype) = R.reshape(
+ scale, R.shape([num_groups, 1])
+ )
+ lv1: R.Tensor((num_groups, channels_per_group),
dtype=dtype) = R.broadcast_to(
+ lv, R.shape([num_groups, channels_per_group])
+ )
+ lv2: R.Tensor((channels,), dtype=dtype) = R.reshape(lv1,
R.shape([channels]))
+ lv3: R.Tensor((num_groups, 1), dtype=dtype) = R.reshape(
+ bias, R.shape([num_groups, 1])
+ )
+ lv4: R.Tensor((num_groups, channels_per_group),
dtype=dtype) = R.broadcast_to(
+ lv3, R.shape([num_groups, channels_per_group])
+ )
+ lv5: R.Tensor((channels,), dtype=dtype) = R.reshape(lv4,
R.shape([channels]))
+ gv: R.Tensor(input_shape, dtype=dtype) = R.nn.group_norm(
+ input,
+ lv2,
+ lv5,
+ num_groups=num_groups,
+ channel_axis=1,
+ axes=axes,
+ epsilon=epsilon,
+ )
+ R.output(gv)
+ return gv
+
+ return ExpectedGroupNormOpset18
+
+ if opset == 21 and stash_type == 1 and dtype != "float32":
+
+ @I.ir_module
+ class ExpectedGroupNormOpset21Stash:
+ @R.function
+ def main(
+ input: R.Tensor(input_shape, dtype=dtype),
+ scale: R.Tensor(scale_shape, dtype=dtype),
+ bias: R.Tensor(bias_shape, dtype=dtype),
+ ) -> R.Tensor(input_shape, dtype=dtype):
+ R.func_attr({"num_input": 3})
+ with R.dataflow():
+ lv: R.Tensor(input_shape, dtype="float32") =
R.astype(input, dtype="float32")
+ lv1: R.Tensor(scale_shape, dtype="float32") =
R.astype(scale, dtype="float32")
+ lv2: R.Tensor(scale_shape, dtype="float32") =
R.ones_like(lv1)
+ lv3: R.Tensor(bias_shape, dtype="float32") =
R.astype(bias, dtype="float32")
+ lv4: R.Tensor(bias_shape, dtype="float32") =
R.zeros_like(lv3)
+ lv5: R.Tensor(input_shape, dtype="float32") =
R.nn.group_norm(
+ lv,
+ lv2,
+ lv4,
+ num_groups=num_groups,
+ channel_axis=1,
+ axes=axes,
+ epsilon=epsilon,
+ center=False,
+ scale=False,
+ )
+ lv6: R.Tensor(input_shape, dtype=dtype) = R.astype(lv5,
dtype=dtype)
+ lv7: R.Tensor(affine_shape, dtype=dtype) = R.reshape(
+ scale, R.shape(affine_shape)
+ )
+ lv8: R.Tensor(input_shape, dtype=dtype) = R.multiply(lv6,
lv7)
+ lv9: R.Tensor(affine_shape, dtype=dtype) = R.reshape(
+ bias, R.shape(affine_shape)
+ )
+ gv: R.Tensor(input_shape, dtype=dtype) = R.add(lv8, lv9)
+ R.output(gv)
+ return gv
+
+ return ExpectedGroupNormOpset21Stash
+
+ if opset == 21:
+
+ @I.ir_module
+ class ExpectedGroupNormOpset21:
+ @R.function
+ def main(
+ input: R.Tensor(input_shape, dtype=dtype),
+ scale: R.Tensor(scale_shape, dtype=dtype),
+ bias: R.Tensor(bias_shape, dtype=dtype),
+ ) -> R.Tensor(input_shape, dtype=dtype):
+ R.func_attr({"num_input": 3})
+ with R.dataflow():
+ lv: R.Tensor(scale_shape, dtype=dtype) = R.ones_like(scale)
+ lv1: R.Tensor(bias_shape, dtype=dtype) = R.zeros_like(bias)
+ lv2: R.Tensor(input_shape, dtype=dtype) = R.nn.group_norm(
+ input,
+ lv,
+ lv1,
+ num_groups=num_groups,
+ channel_axis=1,
+ axes=axes,
+ epsilon=epsilon,
+ center=False,
+ scale=False,
+ )
+ lv3: R.Tensor(affine_shape, dtype=dtype) = R.reshape(
+ scale, R.shape(affine_shape)
+ )
+ lv4: R.Tensor(input_shape, dtype=dtype) = R.multiply(lv2,
lv3)
+ lv5: R.Tensor(affine_shape, dtype=dtype) = R.reshape(
+ bias, R.shape(affine_shape)
+ )
+ gv: R.Tensor(input_shape, dtype=dtype) = R.add(lv4, lv5)
+ R.output(gv)
+ return gv
+
+ return ExpectedGroupNormOpset21
+
+ raise AssertionError(f"No GroupNormalization expected IR for
opset={opset}")
+
+
+def test_group_norm():
+ def verify_group_norm(
+ input_shape: list[int],
+ scale_shape: list[int],
+ bias_shape: list[int],
+ num_groups: int,
+ expected,
+ opset: int = 21,
+ dtype: int = TensorProto.FLOAT,
+ stash_type: int = 1,
+ ):
+ attrs = {"num_groups": num_groups, "epsilon": 1e-5}
+ if opset == 21:
+ attrs["stash_type"] = stash_type
+
+ node = helper.make_node(
+ "GroupNormalization", ["input", "scale", "bias"], ["output"],
**attrs
+ )
+ graph = helper.make_graph(
+ [node],
+ "group_norm_test",
+ inputs=[
+ helper.make_tensor_value_info("input", dtype,
list(input_shape)),
+ helper.make_tensor_value_info("scale", dtype,
list(scale_shape)),
+ helper.make_tensor_value_info("bias", dtype, list(bias_shape)),
+ ],
+ outputs=[helper.make_tensor_value_info("output", dtype,
list(input_shape))],
+ )
+
+ model = helper.make_model(
+ graph,
+ producer_name="group_norm_test",
+ opset_imports=[helper.make_opsetid("", opset)],
+ )
+ tvm_model = from_onnx(model, opset=opset, keep_params_in_input=True)
+ tvm_model["main"] = tvm_model["main"].without_attr("params")
+ expected = tvm.IRModule(expected.functions)
+ for gv in expected.get_global_vars():
+ if gv.name_hint != "main":
+ expected.update_func(gv, tvm_model[gv.name_hint])
+ tvm.ir.assert_structural_equal(tvm_model, expected)
+
+ for input_shape, scale_shape, bias_shape, num_groups, opset, dtype,
dtype_str, stash_type in [
+ ([1, 4, 2, 2], [2], [2], 2, 18, TensorProto.FLOAT, "float32", 1),
+ ([1, 4, 2, 2], [4], [4], 2, 21, TensorProto.FLOAT, "float32", 1),
+ ([1, 4, 8], [4], [4], 2, 21, TensorProto.FLOAT, "float32", 1),
+ ([1, 4, 2, 2], [4], [4], 2, 21, TensorProto.FLOAT16, "float16", 1),
+ ]:
+ verify_group_norm(
+ input_shape,
+ scale_shape,
+ bias_shape,
+ num_groups,
+ _make_group_norm_expected_ir(
+ input_shape,
+ scale_shape,
+ bias_shape,
+ num_groups,
+ opset=opset,
+ dtype=dtype_str,
+ stash_type=stash_type,
+ ),
+ opset=opset,
+ dtype=dtype,
+ stash_type=stash_type,
+ )
+
+ for bad_stash_type in [0, 10, 11, 16]:
+ with pytest.raises(ValueError, match="stash_type=1"):
+ verify_group_norm(
+ [1, 4, 2, 2],
+ [4],
+ [4],
+ 2,
+ _make_group_norm_expected_ir(
+ [1, 4, 2, 2],
+ [4],
+ [4],
+ 2,
+ opset=21,
+ dtype="float16",
+ stash_type=1,
+ ),
+ opset=21,
+ dtype=TensorProto.FLOAT16,
+ stash_type=bad_stash_type,
+ )
+
+ with pytest.raises(ValueError, match="currently only supports float32"):
+ verify_group_norm(
+ [1, 4, 2, 2],
+ [2],
+ [2],
+ 2,
+ _make_group_norm_expected_ir(
+ [1, 4, 2, 2],
+ [2],
+ [2],
+ 2,
+ opset=18,
+ dtype="float16",
+ ),
+ opset=18,
+ dtype=TensorProto.FLOAT16,
+ )
+
+
# TODO Enable dynamism
@pytest.mark.parametrize("dynamic", [False])
def test_skiplayernormalization(dynamic):
