[Fix][Relax][ONNX] Cast BatchNorm params to input dtype

[viiccwen]

Fixes #19977.

ONNX BatchNormalization allows the input/output tensor dtype, scale/bias dtype, and mean/variance dtype to be separate floating-point type parameters.

For example, a valid ONNX model may use float16 data with float32 gamma, beta, mean, and variance tensors.

The Relax batch_norm operator currently requires all five input tensors to have the same dtype. The ONNX frontend previously forwarded the ONNX inputs directly to relax.nn.batch_norm, causing import to fail during normalization
for mixed-dtype ONNX models.

This patch casts the ONNX BatchNormalization parameter tensors (scale, bias, mean, and var) to the data tensor dtype before calling Relax batch_norm.

This preserves the ONNX output dtype, which follows the input data dtype, while keeping the fix localized to the frontend compatibility layer.

The regression test builds a minimal ONNX BatchNormalization graph with float16 data and float32 parameters, imports it through the Relax ONNX frontend, and checks that the generated Relax batch_norm call receives same-dtype inputs.

Verification:

• python -m pytest tests/python/relax/test_frontend_onnx.py::test_batch_norm_mixed_dtype_params tests/python/relax/test_frontend_onnx.py::test_batch_norm_defaults_to_inference_mode -q

[gemini-code-assist]

Code Review

This pull request updates the ONNX frontend's BatchNormalization implementation to automatically cast the scale, bias, mean, and variance parameters to match the input data type, preventing type mismatches. A new unit test is also added to verify this mixed-precision behavior. The review feedback suggests optimizing this casting logic for constant parameters by directly modifying the underlying numpy arrays and recreating relax.Constant nodes, thereby avoiding unnecessary astype operations in the Relax AST.

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.

[tlopex]

Thanks for fixing! But could you avoid casting the float32 parameters down to the data dtype, as this can lose precision and change the result.

Since Relax batch_norm requires matching dtypes, please promote all five inputs to a common computation dtype (float32 for this case), run batch_norm, and cast Y back to the original data dtype.

The test can verify that the batch_norm inputs are float32 while the final output remains float16 I think

[viiccwen]

Thanks for fixing! But could you avoid casting the float32 parameters down to the data dtype, as this can lose precision and change the result.

Since Relax batch_norm requires matching dtypes, please promote all five inputs to a common computation dtype (float32 for this case), run batch_norm, and cast Y back to the original data dtype.

The test can verify that the batch_norm inputs are float32 while the final output remains float16 I think

Thx for pointing out! I'll fix it later. : D

[viiccwen]

Hello @tlopex,

Updated the implementation. (Follow by ONNX BatchNormalization-15 schema)

The strategy is:

• Preserve ONNX's dtype groups: X/Y use T, scale/bias use T1, and mean/variance outputs use T2.
• For the reported mixed-precision case (float16 X with float16/float32 parameters and statistics), run Relax batch_norm in float32.
• Cast Y back to T and the running mean/variance back to T2.
• In training mode, float16 inputs also use float32 computation, following the ONNX requirement to compute training statistics in float to avoid overflow.
• Other mixed-dtype combinations are rejected explicitly instead of applying an unspecified general promotion policy. (can be discussed how to handle. : D)

[tlopex]

LGTM