viiccwen opened a new pull request, #19979: URL: https://github.com/apache/tvm/pull/19979
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` ``` -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
