[Bug] LegalizeOps fails to process LayerNormalization of valid onnx model: tvm.error.InternalError: Check failed: shape.size() == indices.size() (1 vs. 0)

[coffezhou]

Expected behavior

LegalizeOps should process the LayerNormalization of valid onnx model correctly.

Actual behavior

The onnx model can be run by onnxruntime and outputs the following results:

 [array([0., 0., 0., 0., 0.], dtype=float32)]

However, when legalize the LayerNormalization, TVM crashes as follows:

Traceback (most recent call last):
  File "/home/carla/Documents/test/test.py", line 41, in <module>
    main()
  File "/home/carla/Documents/test/test.py", line 37, in main
    tvm_model = relax.transform.LegalizeOps()(tvm_model)
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/carla/Documents/tvm/python/tvm/ir/transform.py", line 238, in __call__
    return _ffi_transform_api.RunPass(self, mod)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "tvm/ffi/cython/./function.pxi", line 228, in tvm.ffi.core.Function.__call__
  File "tvm/ffi/cython/./function.pxi", line 281, in tvm.ffi.core.tvm_ffi_callback
  File "/home/carla/Documents/tvm/python/tvm/relax/transform/legalize_ops/nn.py", line 613, in _nn_layer_norm
    return bb.call_te(
  File "/home/carla/Documents/tvm/python/tvm/relax/block_builder.py", line 356, in call_te
    tir_func, call_args, output_sinfo, tir_vars = gen_call_tir_inputs(func, *args, **kwargs)
  File "/home/carla/Documents/tvm/python/tvm/relax/utils.py", line 354, in gen_call_tir_inputs
    te_out = func(*te_args, **te_kwargs)
  File "/home/carla/Documents/tvm/python/tvm/topi/nn/layer_norm.py", line 48, in layer_norm
    return cpp.nn.layer_norm(data, gamma, beta, axis, epsilon)
  File "tvm/ffi/cython/./function.pxi", line 228, in tvm.ffi.core.Function.__call__
tvm.error.InternalError: Check failed: shape.size() == indices.size() (1 vs. 0) : Tensor dimension mismatch in read ndim = 1, indices.size=0

Environment

OS: Ubuntu 20.04 TVM: 0.21.dev0 (3db71bb3a)

Steps to reproduce

This bug can be reproduced by the following code with the model in the attachment. As shown in the code, the model can be executed by onnxruntime.

import sys

import numpy as np
import onnx
import onnxruntime

import tvm
from tvm import relax
from tvm.relax.frontend.onnx import from_onnx

import pickle

            
def main():
    onnx_model = onnx.load("a552.onnx")
    
    shape_onnx_model = onnx.shape_inference.infer_shapes(onnx_model)
    onnx.save(shape_onnx_model, '1111.onnx')
    
    with open("inputs.pkl", "rb") as fp:
        inputs = pickle.load(fp)
    
    try:
        ort_session = onnxruntime.InferenceSession(
            onnx_model.SerializeToString(), providers=["CPUExecutionProvider"]
        )
        ort_output = ort_session.run([], inputs)
    except Exception as e:
        print(e)
        sys.exit(1)
    print("ONNXRuntime:\n", ort_output)   
    # Convert the onnx model into relax through the onnx importer.
    tvm_model = from_onnx(onnx_model, keep_params_in_input=True)
    # Convert operators for inference mode.
    tvm_model = relax.transform.DecomposeOpsForInference()(tvm_model)
    # Legalize any relax ops into tensorir.
    tvm_model = relax.transform.LegalizeOps()(tvm_model)
    
if __name__ == "__main__":
    
    main()

testcase.zip

Triage

• needs-triage