[Bug] Relax ONNX frontend fails to import mixed-dtype BatchNormalization

[Shirley4042]

Expected behavior

The ONNX model should be imported successfully by the Relax ONNX frontend.

Actual behavior

from_onnx fails while converting the BatchNormalization node because
relax.nn.batch_norm requires all inputs to have the same dtype.

tvm.error.InternalError:
relax.nn.batch_norm requires all the input tensors to have the same dtype.
However, the gamma has dtype float32, which is different from the input
data's dtype float16.

Environment

TVM: 0.25.dev0, built from source
Python: 3.10
OS: Ubuntu 22.04
Target: llvm (CPU)
ONNX frontend: tvm.relax.frontend.onnx.from_onnx

Steps to reproduce

import numpy as np
import onnx
from onnx import TensorProto, helper, numpy_helper
from tvm.relax.frontend.onnx import from_onnx
data = helper.make_tensor_value_info(
    "data",
    TensorProto.FLOAT16,
    [1, 3, 2, 2],
)
output = helper.make_tensor_value_info(
    "output",
    TensorProto.FLOAT16,
    [1, 3, 2, 2],
)
params = [
    numpy_helper.from_array(
        np.array([1.0, 1.5, 2.0], dtype=np.float32),
        name="gamma",
    ),
    numpy_helper.from_array(
        np.array([0.0, 0.1, -0.1], dtype=np.float32),
        name="beta",
    ),
    numpy_helper.from_array(
        np.array([0.2, -0.3, 0.4], dtype=np.float32),
        name="mean",
    ),
    numpy_helper.from_array(
        np.array([1.0, 1.5, 2.0], dtype=np.float32),
        name="var",
    ),
]
node = helper.make_node(
    "BatchNormalization",
    inputs=["data", "gamma", "beta", "mean", "var"],
    outputs=["output"],
    epsilon=1e-5,
    momentum=0.9,
    training_mode=0,
)
graph = helper.make_graph(
    [node],
    "mixed_dtype_batchnorm",
    [data],
    [output],
    initializer=params,
)
model = helper.make_model(
    graph,
    opset_imports=[helper.make_opsetid("", 15)],
)
onnx.checker.check_model(model, full_check=True)
mod = from_onnx(
    model,
    keep_params_in_input=False,
)

Analysis

The Relax ONNX frontend directly converts the ONNX BatchNormalization node into relax.nn.batch_norm.
The ONNX model contains mixed floating-point dtypes:

data:float16, gamma:float32, beta:float32, moving_mean: float32, moving_var: float32

The Relax batch_norm operator currently requires all five inputs to have the same dtype. Therefore, the generated Relax expression fails during BlockBuilder.normalize() before the ONNX model can be imported.

[viiccwen]

Hello @Shirley4042, I'll work on this. : )