apache--tvm/tests/python/relax/test_frontend_onnx.py

# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.
# pylint: disable=unused-argument
# ruff: noqa: E501, F841
"""
ONNX testcases
================
This file is a test script to test Relax ONNX frontend coverage.
"""

# Allow TVMScript expected IR to capture shape/dtype names used only in annotations.
from __future__ import annotations

from typing import Literal

import numpy as np
import pytest

pytest.importorskip("onnx")
pytest.importorskip("onnxruntime")

import onnx
import onnxruntime
import tvm_ffi
from onnx import ModelProto, TensorProto, helper, numpy_helper

import tvm
import tvm.testing
from tvm import relax
from tvm.relax.frontend.onnx import from_onnx
from tvm.script import ir as I
from tvm.script import relax as R
from tvm.script import tirx as T

bg = np.random.MT19937(0)
rg = np.random.Generator(bg)


def collect_relax_call_ops(func):
    call_ops = set()

    def _visit(expr):
        if isinstance(expr, relax.Call) and isinstance(expr.op, tvm.ir.Op):
            call_ops.add(expr.op.name)

    relax.analysis.post_order_visit(func.body, _visit)
    return call_ops


def generate_random_inputs(
    model: ModelProto, inputs: dict[str, np.ndarray] | None = None
) -> dict[str, np.ndarray]:
    input_values = {}
    # Iterate through model inputs and extract their shape.
    for i in model.graph.input:
        if inputs is not None and i.name in inputs and inputs[i.name] is not None:
            input_values[i.name] = inputs[i.name]
            continue
        shape = []
        for dim in i.type.tensor_type.shape.dim:
            shape.append(dim.dim_value)

        input_values[i.name] = generate_random_value(shape, i.type.tensor_type.elem_type)

    return input_values


def generate_random_value(shape, elem_type) -> np.ndarray:
    # Extract datatype for the input.
    if elem_type:
        dtype = str(helper.tensor_dtype_to_np_dtype(elem_type))
    else:
        dtype = "float32"

    # Generate random inputs for each input.
    if dtype == "bool":
        # random_value = np.random.choice(a=[False, True], size=shape)
        random_value = rg.choice(a=[False, True], size=shape)
    elif dtype.startswith("int"):
        # Keep non-zero values
        random_value = rg.integers(low=-63, high=63, size=shape).astype(dtype)
        random_value[random_value <= 0] -= 1
    else:
        random_value = rg.standard_normal(size=shape).astype(dtype)

    return random_value


def check_correctness(
    model: ModelProto,
    inputs: dict[str, np.ndarray] | None = None,
    ir_version: int = 8,
    opset: int = 14,
    rtol: float = 1e-7,
    atol: float = 1e-5,
    check_dtypes: bool = False,
) -> None:
    """Run an onnx model in both onnxruntime and TVM through our importer
       confirm that the results match. Otherwise, an exception will be raised.

    Parameters
    ----------
    model: ModelProto
        The input onnx model that should be tested.
    inputs: Optional[Dict[str, np.ndarray]]
        An optional dictionary containing values for each input in the onnx model.
    ir_version: int
        Which version of the onnx IR to use.
    opset: int
        The opset version to use for the onnx importer.
    atol: float
        Set the tolerance of correctness checking. Some ops may be show more
        arithmetic variance than others.
    check_dtypes: bool
        Check if data types are the same.
    """
    # Configure model format.
    if ir_version is not None:
        model.ir_version = ir_version
    if opset is not None:
        model.opset_import[0].version = opset

    # If inputs are not provided, extract them from the onnx graph and produce random
    # values that we'll use for testing.
    inputs = generate_random_inputs(model, inputs)

    # Run the model through onnx to get the expected result.
    ort_session = onnxruntime.InferenceSession(
        model.SerializeToString(), providers=["CPUExecutionProvider"]
    )
    ort_output = ort_session.run([], inputs)

    # Convert the onnx model into relax through the onnx importer.
    tvm_model = from_onnx(model, opset=opset, 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)

    # Separate model from parameters.
    tvm_model, params = relax.frontend.detach_params(tvm_model)
    # Compile the relax graph into a VM then run.
    with tvm.transform.PassContext(opt_level=3):
        ex = tvm.compile(tvm_model, target="llvm")
        vm = relax.VirtualMachine(ex, tvm.cpu())
    # Prepare inputs.
    input_list = [
        inputs[key.name_hint] for key in tvm_model["main"].params if key.name_hint in inputs
    ]
    if params:
        input_list += params["main"]

    # Run model and check outputs.
    vm.set_input("main", *input_list)
    vm.invoke_stateful("main")
    tvm_output = vm.get_outputs("main")
    # Wrap as a list if there is only one output.
    if len(ort_output) == 1:
        # Do not check the output number for TVM
        # As for sequence output, the TVM output is a Tuple
        # while the ONNX output number is one, which is a list
        tvm_output = [tvm_output]

    def _get_numpy_subdtype(narray):
        if np.issubdtype(narray.dtype, np.integer):
            return "integer"
        elif np.issubdtype(narray.dtype, np.floating):
            return "floating"
        elif np.issubdtype(narray.dtype, np.bool_):
            return "bool"
        elif np.issubdtype(narray.dtype, np.complexfloating):
            return "complexfloating"
        else:
            return "other"

    def _check_output(tvm_out, ort_out):
        if isinstance(tvm_out, tuple) and isinstance(ort_out, tvm_ffi.Shape | list):
            assert len(tvm_out) == len(ort_out), "Unequal number of outputs"
            for tvm_out_i, ort_out_i in zip(tvm_out, ort_out):
                _check_output(tvm_out_i, ort_out_i)
        elif isinstance(tvm_out, tvm.runtime.Tensor) and isinstance(ort_out, np.ndarray):
            if check_dtypes:
                assert tvm_out.numpy().dtype == ort_out.dtype
            tvm.testing.assert_allclose(tvm_out.numpy(), ort_out, rtol=rtol, atol=atol)
        elif isinstance(tvm_out, tvm_ffi.Shape) and isinstance(ort_out, np.ndarray):
            shape_out = tvm.runtime.tensor([int(i) for i in tvm_out])
            if check_dtypes:
                assert _get_numpy_subdtype(shape_out.numpy()) == _get_numpy_subdtype(ort_out)
            tvm.testing.assert_allclose(shape_out.numpy(), ort_out, rtol=rtol, atol=atol)
        elif isinstance(tvm_out, int | float | bool) and isinstance(ort_out, np.ndarray):
            if check_dtypes:
                assert _get_numpy_subdtype(np.array(tvm_out)) == _get_numpy_subdtype(ort_out)
            tvm.testing.assert_allclose(np.array(tvm_out), ort_out, rtol=rtol, atol=atol)
        else:
            raise ValueError(f"Unsupported types: {type(tvm_out)}, {type(ort_out)}")

    # Check that number of outputs match.
    assert len(tvm_output) == len(ort_output), "Unequal number of outputs"
    for tvm_out, ort_out in zip(tvm_output, ort_output):
        # TODO Allow configurable tolerance.
        if ort_out is not None:
            _check_output(tvm_out, ort_out)


def run_in_tvm(
    model: ModelProto,
    inputs: dict[str, np.ndarray] | None = None,
    ir_version: int = 8,
    opset: int = 14,
):
    if ir_version is not None:
        model.ir_version = ir_version
    if opset is not None:
        for opset_import in model.opset_import:
            if opset_import.domain in ["", "ai.onnx"]:
                opset_import.version = opset
                break

    inputs = generate_random_inputs(model, inputs)
    tvm_model = from_onnx(model, opset=opset, keep_params_in_input=True)
    tvm_model = relax.transform.DecomposeOpsForInference()(tvm_model)
    tvm_model = relax.transform.LegalizeOps()(tvm_model)
    tvm_model, params = relax.frontend.detach_params(tvm_model)

    with tvm.transform.PassContext(opt_level=3):
        ex = tvm.compile(tvm_model, target="llvm")
        vm = relax.VirtualMachine(ex, tvm.cpu())

    input_list = [
        inputs[key.name_hint] for key in tvm_model["main"].params if key.name_hint in inputs
    ]
    if params:
        input_list += params["main"]

    vm.set_input("main", *input_list)
    vm.invoke_stateful("main")
    return vm.get_outputs("main")


@pytest.mark.parametrize(
    "input_names, expected_names",
    [
        ([".", "123"], ["_", "input_123"]),
        ([".", "_"], ["_", "__1"]),
        (["123", "input_123"], ["input_123", "input_123_1"]),
    ],
)
def test_sanitize(input_names, expected_names):
    node = helper.make_node("Add", inputs=input_names, outputs=["output"])
    graph = helper.make_graph(
        [node],
        "test",
        inputs=[
            helper.make_tensor_value_info(str(var), TensorProto.FLOAT, [32, 32])
            for var in input_names
        ],
        outputs=[
            helper.make_tensor_value_info("output", TensorProto.FLOAT, [32, 32]),
        ],
    )
    model = helper.make_model(graph, producer_name="test_sanitizer")

    tvm_model = from_onnx(model)

    for i, param in enumerate(tvm_model["main"].params):
        assert param.name_hint == expected_names[i]


def verify_unary(
    op_name,
    shape,
    attrs={},
    domain=None,
    input_dtype=TensorProto.FLOAT,
    output_dtype=TensorProto.FLOAT,
    opset=14,
):
    test_node = helper.make_node(op_name, ["x"], ["y"], **attrs, domain=domain)
    graph = helper.make_graph(
        [test_node],
        "elemwise_test",
        inputs=[
            helper.make_tensor_value_info("x", input_dtype, shape),
        ],
        outputs=[helper.make_tensor_value_info("y", output_dtype, shape)],
    )

    model = helper.make_model(graph, producer_name="elemwise_test")
    check_correctness(model, opset=opset)


def make_unary_model(
    op_name,
    shape,
    attrs=None,
    domain=None,
    input_dtype=TensorProto.FLOAT,
    output_dtype=TensorProto.FLOAT,
):
    attrs = attrs or {}
    test_node = helper.make_node(op_name, ["x"], ["y"], **attrs, domain=domain)
    graph = helper.make_graph(
        [test_node],
        "elemwise_structural_test",
        inputs=[
            helper.make_tensor_value_info("x", input_dtype, shape),
        ],
        outputs=[helper.make_tensor_value_info("y", output_dtype, shape)],
    )
    return helper.make_model(graph, producer_name="elemwise_structural_test")


def verify_binary(
    op_name, shape_a, shape_b, shape_c, attrs={}, domain=None, dtype=TensorProto.FLOAT, opset=14
):
    test_node = helper.make_node(op_name, ["a", "b"], ["c"], **attrs, domain=domain)
    graph = helper.make_graph(
        [test_node],
        "binary_test",
        inputs=[
            helper.make_tensor_value_info("a", dtype, shape_a),
            helper.make_tensor_value_info("b", dtype, shape_b),
        ],
        outputs=[helper.make_tensor_value_info("c", dtype, shape_c)],
    )

    model = helper.make_model(graph, producer_name="binary_test")
    check_correctness(model, opset=opset, check_dtypes=True)


def verify_binary_scalar(op_name, attrs={}, domain=None, dtype=TensorProto.INT32, opset=14):
    a = make_constant_node("a", dtype, [], [4])
    b = make_constant_node("b", dtype, [], [8])
    test_node = helper.make_node(op_name, ["a", "b"], ["c"], **attrs, domain=domain)
    graph = helper.make_graph(
        [a, b, test_node],
        "binary_test",
        inputs=[],
        outputs=[helper.make_tensor_value_info("c", dtype, ())],
    )

    model = helper.make_model(graph, producer_name="binary_test")
    model.opset_import[0].version = opset
    tvm_model = from_onnx(model, opset=opset, keep_params_in_input=True)

    dtype_str = str(helper.tensor_dtype_to_np_dtype(dtype))
    lhs = np.array(4, dtype=dtype_str)
    rhs = np.array(8, dtype=dtype_str)
    op = {
        "Add": np.add,
        "Sub": np.subtract,
        "Mul": np.multiply,
        "Div": np.divide,
        "Pow": np.power,
        "Mod": np.mod if attrs.get("fmod", 0) else np.fmod,
    }[op_name]
    expected_value = op(lhs, rhs).astype(dtype_str)

    @I.ir_module
    class Expected:
        @R.function
        def main() -> R.Tensor((), dtype=dtype_str):
            R.func_attr({"num_input": 0})
            with R.dataflow():
                gv: R.Tensor((), dtype=dtype_str) = R.const(expected_value.item(), dtype_str)
                R.output(gv)
            return gv

    tvm.ir.assert_structural_equal(tvm_model, Expected)


def verify_compare(op_name, shape, attrs={}, domain=None):
    test_node = helper.make_node(op_name, ["a", "b"], ["c"], **attrs, domain=domain)
    graph = helper.make_graph(
        [test_node],
        "compare_test",
        inputs=[
            helper.make_tensor_value_info("a", TensorProto.FLOAT, shape),
            helper.make_tensor_value_info("b", TensorProto.FLOAT, shape),
        ],
        outputs=[helper.make_tensor_value_info("c", TensorProto.BOOL, shape)],
    )

    model = helper.make_model(graph, producer_name="compare_test")
    check_correctness(model)


@pytest.mark.parametrize("dynamic", [True, False])
def test_matmul(dynamic):
    matmul_node = helper.make_node("MatMul", ["a", "b"], ["c"])

    a_shape = [32, 48]
    b_shape = [48, 64]
    output_shape = [32, 64]

    if dynamic:
        a_shape = ["?", "?"]

    graph = helper.make_graph(
        [matmul_node],
        "matmul_test",
        inputs=[
            helper.make_tensor_value_info("a", TensorProto.FLOAT, a_shape),
        ],
        initializer=[
            helper.make_tensor(
                "b", TensorProto.FLOAT, b_shape, np.random.normal(size=b_shape).astype("float32")
            )
        ],
        outputs=[helper.make_tensor_value_info("c", TensorProto.FLOAT, output_shape)],
    )

    model = helper.make_model(graph, producer_name="matmul_test")
    inputs = None
    if dynamic:
        inputs = {
            "a": np.random.normal(size=[32, 48]).astype("float32"),
        }
    check_correctness(model, inputs)


def test_matmulinteger16():
    def verify_matmulinteger16(a_dtype, b_dtype, a_shape, b_shape, expected):
        out_dtype = np.uint32 if a_dtype == np.uint16 and b_dtype == np.uint16 else np.int32
        output_shape = [
            *np.broadcast_shapes(tuple(a_shape[:-2]), tuple(b_shape[:-2])),
            a_shape[-2],
            b_shape[-1],
        ]

        node = helper.make_node("MatMulInteger16", ["a", "b"], ["y"], domain="com.microsoft")
        graph = helper.make_graph(
            [node],
            "matmulinteger16_test",
            inputs=[
                helper.make_tensor_value_info(
                    "a", helper.np_dtype_to_tensor_dtype(np.dtype(a_dtype)), a_shape
                ),
                helper.make_tensor_value_info(
                    "b", helper.np_dtype_to_tensor_dtype(np.dtype(b_dtype)), b_shape
                ),
            ],
            outputs=[
                helper.make_tensor_value_info(
                    "y",
                    helper.np_dtype_to_tensor_dtype(np.dtype(out_dtype)),
                    output_shape,
                )
            ],
        )
        model = helper.make_model(
            graph,
            producer_name="matmulinteger16_test",
            opset_imports=[helper.make_opsetid("", 18), helper.make_opsetid("com.microsoft", 1)],
        )
        model.ir_version = 11

        tvm_model = from_onnx(model, opset=18, keep_params_in_input=True)
        tvm.ir.assert_structural_equal(tvm_model, expected)

    @I.ir_module
    class ExpectedInt16:
        @R.function
        def main(
            a: R.Tensor((2, 3), dtype="int16"),
            b: R.Tensor((3, 4), dtype="int16"),
        ) -> R.Tensor((2, 4), dtype="int32"):
            R.func_attr({"num_input": 2})
            with R.dataflow():
                lv: R.Tensor((2, 3), dtype="int32") = R.astype(a, dtype="int32")
                lv1: R.Tensor((3, 4), dtype="int32") = R.astype(b, dtype="int32")
                gv: R.Tensor((2, 4), dtype="int32") = R.matmul(lv, lv1)
                R.output(gv)
            return gv

    @I.ir_module
    class ExpectedUInt16:
        @R.function
        def main(
            a: R.Tensor((2, 3), dtype="uint16"),
            b: R.Tensor((3, 4), dtype="uint16"),
        ) -> R.Tensor((2, 4), dtype="uint32"):
            R.func_attr({"num_input": 2})
            with R.dataflow():
                lv: R.Tensor((2, 3), dtype="uint32") = R.astype(a, dtype="uint32")
                lv1: R.Tensor((3, 4), dtype="uint32") = R.astype(b, dtype="uint32")
                gv: R.Tensor((2, 4), dtype="uint32") = R.matmul(lv, lv1)
                R.output(gv)
            return gv

    @I.ir_module
    class ExpectedMixedBatched:
        @R.function
        def main(
            a: R.Tensor((2, 1, 3, 5), dtype="int16"),
            b: R.Tensor((1, 2, 5, 4), dtype="uint16"),
        ) -> R.Tensor((2, 2, 3, 4), dtype="int32"):
            R.func_attr({"num_input": 2})
            with R.dataflow():
                lv: R.Tensor((2, 1, 3, 5), dtype="int32") = R.astype(a, dtype="int32")
                lv1: R.Tensor((1, 2, 5, 4), dtype="int32") = R.astype(b, dtype="int32")
                gv: R.Tensor((2, 2, 3, 4), dtype="int32") = R.matmul(lv, lv1)
                R.output(gv)
            return gv

    verify_matmulinteger16(np.int16, np.int16, [2, 3], [3, 4], ExpectedInt16)
    verify_matmulinteger16(np.uint16, np.uint16, [2, 3], [3, 4], ExpectedUInt16)
    verify_matmulinteger16(
        np.int16,
        np.uint16,
        [2, 1, 3, 5],
        [1, 2, 5, 4],
        ExpectedMixedBatched,
    )


def test_matmulinteger16_ir():
    node = helper.make_node("MatMulInteger16", ["a", "b"], ["y"], domain="com.microsoft")
    graph = helper.make_graph(
        [node],
        "matmulinteger16_ir_test",
        inputs=[
            helper.make_tensor_value_info("a", TensorProto.UINT16, [2, 3]),
            helper.make_tensor_value_info("b", TensorProto.UINT16, [3, 4]),
        ],
        outputs=[helper.make_tensor_value_info("y", TensorProto.UINT32, [2, 4])],
    )
    model = helper.make_model(
        graph,
        producer_name="matmulinteger16_ir_test",
        opset_imports=[helper.make_opsetid("", 18), helper.make_opsetid("com.microsoft", 1)],
    )
    model.ir_version = 11

    tvm_model = from_onnx(model, opset=18, keep_params_in_input=True)

    @I.ir_module
    class Expected:
        @R.function
        def main(
            a: R.Tensor((2, 3), dtype="uint16"),
            b: R.Tensor((3, 4), dtype="uint16"),
        ) -> R.Tensor((2, 4), dtype="uint32"):
            R.func_attr({"num_input": 2})
            with R.dataflow():
                lv: R.Tensor((2, 3), dtype="uint32") = R.astype(a, dtype="uint32")
                lv1: R.Tensor((3, 4), dtype="uint32") = R.astype(b, dtype="uint32")
                gv: R.Tensor((2, 4), dtype="uint32") = R.matmul(lv, lv1)
                R.output(gv)
            return gv

    tvm.ir.assert_structural_equal(tvm_model, Expected)


def test_matmulinteger16_invalid_dtype_raises():
    node = helper.make_node("MatMulInteger16", ["a", "b"], ["y"], domain="com.microsoft")
    graph = helper.make_graph(
        [node],
        "matmulinteger16_invalid_dtype_test",
        inputs=[
            helper.make_tensor_value_info("a", TensorProto.INT8, [2, 3]),
            helper.make_tensor_value_info("b", TensorProto.UINT16, [3, 4]),
        ],
        outputs=[helper.make_tensor_value_info("y", TensorProto.INT32, [2, 4])],
    )
    model = helper.make_model(
        graph,
        producer_name="matmulinteger16_invalid_dtype_test",
        opset_imports=[helper.make_opsetid("", 18), helper.make_opsetid("com.microsoft", 1)],
    )
    model.ir_version = 11

    with pytest.raises(ValueError, match="input A"):
        from_onnx(model, opset=18, keep_params_in_input=True)


def test_concat():
    verify_binary("Concat", [1, 32], [1, 32], [2, 32], attrs={"axis": 0})


def test_concat_with_param_shape_value():
    """Concat must handle a 1D-int64 initializer mixed with a ShapeExpr when
    keep_params_in_input=True. Standard pattern in PyTorch-exported ONNX
    models for dynamic-batch Reshape: Reshape(x, Concat(Shape(x)[:1], [12]))."""
    inp = helper.make_tensor_value_info("x", TensorProto.FLOAT, ["N", 3, 4])
    out = helper.make_tensor_value_info("y", TensorProto.FLOAT, ["N", 12])
    twelve = numpy_helper.from_array(np.array([12], dtype=np.int64), "twelve")
    starts = numpy_helper.from_array(np.array([0], dtype=np.int64), "starts")
    ends = numpy_helper.from_array(np.array([1], dtype=np.int64), "ends")
    nodes = [
        helper.make_node("Shape", ["x"], ["x_shape"]),
        helper.make_node("Slice", ["x_shape", "starts", "ends"], ["dyn_n"]),
        helper.make_node("Concat", ["dyn_n", "twelve"], ["new_shape"], axis=0),
        helper.make_node("Reshape", ["x", "new_shape"], ["y"]),
    ]
    graph = helper.make_graph(
        nodes,
        "concat_param_shape",
        [inp],
        [out],
        initializer=[twelve, starts, ends],
    )
    model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 13)])
    model.ir_version = 8
    onnx.checker.check_model(model)
    # Both modes should succeed; previously True crashed with
    # "Op(relax.concat) expects the input to be a Tuple of Tensors".
    from_onnx(model, keep_params_in_input=False)
    from_onnx(model, keep_params_in_input=True)


def test_concat_with_param_tensor_keeps_runtime_param():
    """Concat(input, weight) under keep_params_in_input=True must keep `weight`
    as a runtime param, not fold it into a constant."""
    weight_np = np.arange(8, dtype=np.float32).reshape(2, 4)
    graph = helper.make_graph(
        [helper.make_node("Concat", ["x", "w"], ["y"], axis=0)],
        "concat_param_tensor",
        [helper.make_tensor_value_info("x", TensorProto.FLOAT, [2, 4])],
        [helper.make_tensor_value_info("y", TensorProto.FLOAT, [4, 4])],
        initializer=[numpy_helper.from_array(weight_np, "w")],
    )
    model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 13)])
    model.ir_version = 8
    onnx.checker.check_model(model)

    mod, params = relax.frontend.detach_params(from_onnx(model, keep_params_in_input=True))
    assert "w" in [p.name_hint for p in mod["main"].params]
    assert len(params["main"]) == 1
    np.testing.assert_array_equal(params["main"][0].numpy(), weight_np)


@pytest.mark.parametrize("op_name", ["Add", "Sub", "Mul", "Div", "Pow"])
def test_binary(op_name: str):
    verify_binary(op_name, [1, 32], [1, 32], [1, 32])
    verify_binary_scalar(op_name)


def test_div_integer_constant_zero_divisor_raises_valueerror():
    b_init = numpy_helper.from_array(np.array([3, 0, -2, 1], dtype=np.int32), name="b")
    node = helper.make_node("Div", ["a", "b"], ["y"])
    graph = helper.make_graph(
        [node],
        "div_const_zero",
        [helper.make_tensor_value_info("a", TensorProto.INT32, [4])],
        [helper.make_tensor_value_info("y", TensorProto.INT32, [4])],
        initializer=[b_init],
    )
    model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 18)])
    model.ir_version = 9

    with pytest.raises(
        ValueError, match="ONNX Div with integer inputs encountered divisor value 0"
    ):
        from_onnx(model, opset=18, keep_params_in_input=False)


@pytest.mark.parametrize("int_mode", [True, False])
def test_mod(int_mode: bool):
    if int_mode:
        dtype, fmod = TensorProto.INT32, 0
    else:
        dtype, fmod = TensorProto.FLOAT, 1
    verify_binary("Mod", [1, 32], [1, 32], [1, 32], attrs={"fmod": fmod}, dtype=dtype)
    verify_binary_scalar("Mod", attrs={"fmod": fmod}, dtype=dtype)


SHAPE_PARAMS = [
    ([[32, 32], [32, 32]], [32, 32]),
    ([[32, 1], [1, 2]], [32, 2]),
    (
        [
            [
                32,
            ],
            [
                1,
            ],
        ],
        [
            32,
        ],
    ),
    ([[32, 32, 1, 1], [1, 32, 32]], [32, 32, 32, 32]),
    (
        [
            [32, 32, 1, 1],
            [1, 32, 1],
            [
                32,
            ],
        ],
        [32, 32, 32, 32],
    ),
]


def test_multi_input_broadcasting():
    """Multi-input reductions should import broadcast + stack + reduce."""

    def verify_multi_input_broadcasting(op_name, input_shapes, expected_output_shape, expected):
        num_inputs = len(input_shapes)
        input_names = [f"i{i}" for i in range(num_inputs)]

        input_values_info = []
        for name, shape in zip(input_names, input_shapes):
            input_values_info.append(helper.make_tensor_value_info(name, TensorProto.FLOAT, shape))
        test_node = helper.make_node(op_name, input_names, ["output"])
        output_info = helper.make_tensor_value_info(
            "output", TensorProto.FLOAT, expected_output_shape
        )
        graph = helper.make_graph(
            [test_node],
            f"multi_input_{op_name}_test",
            inputs=input_values_info,
            outputs=[output_info],
        )
        model = helper.make_model(graph, producer_name="multi_input_test")
        tvm_model = from_onnx(model, keep_params_in_input=True)
        tvm.ir.assert_structural_equal(tvm_model, expected)

    def make_expected(op_name, input_shapes, output_shape):
        input_shapes = [tuple(shape) for shape in input_shapes]
        output_shape = tuple(output_shape)
        reduce_op = {
            "Min": R.min,
            "Max": R.max,
            "Sum": R.sum,
            "Mean": R.mean,
        }[op_name]
        input_shape_0 = input_shapes[0]
        input_shape_1 = input_shapes[1]

        if len(input_shapes) == 2:

            @I.ir_module
            class ExpectedMultiInputReduction2:
                @R.function
                def main(
                    i0: R.Tensor(input_shape_0, dtype="float32"),
                    i1: R.Tensor(input_shape_1, dtype="float32"),
                ) -> R.Tensor(output_shape, dtype="float32"):
                    R.func_attr({"num_input": 2})
                    with R.dataflow():
                        lv = R.broadcast_to(i0, R.shape(output_shape))
                        lv1 = R.broadcast_to(i1, R.shape(output_shape))
                        lv2 = R.stack((lv, lv1), axis=0)
                        gv = reduce_op(lv2, axis=[0], keepdims=False)
                        R.output(gv)
                    return gv

            return ExpectedMultiInputReduction2

        input_shape_2 = input_shapes[2]

        @I.ir_module
        class ExpectedMultiInputReduction3:
            @R.function
            def main(
                i0: R.Tensor(input_shape_0, dtype="float32"),
                i1: R.Tensor(input_shape_1, dtype="float32"),
                i2: R.Tensor(input_shape_2, dtype="float32"),
            ) -> R.Tensor(output_shape, dtype="float32"):
                R.func_attr({"num_input": 3})
                with R.dataflow():
                    lv = R.broadcast_to(i0, R.shape(output_shape))
                    lv1 = R.broadcast_to(i1, R.shape(output_shape))
                    lv2 = R.broadcast_to(i2, R.shape(output_shape))
                    lv3 = R.stack((lv, lv1, lv2), axis=0)
                    gv = reduce_op(lv3, axis=[0], keepdims=False)
                    R.output(gv)
                return gv

        return ExpectedMultiInputReduction3

    for input_shapes, output_shape in SHAPE_PARAMS:
        for op_name in ["Min", "Max", "Sum", "Mean"]:
            verify_multi_input_broadcasting(
                op_name,
                input_shapes,
                output_shape,
                make_expected(op_name, input_shapes, output_shape),
            )


@pytest.mark.parametrize("op_name", ["Less", "LessOrEqual", "Greater", "GreaterOrEqual"])
def test_compare(op_name: str):
    verify_compare(op_name, [1, 32])


@pytest.mark.parametrize("op_name", ["And", "Or", "Xor"])
def test_binary_bool(op_name: str):
    verify_binary(op_name, [32, 32], [32, 32], [32, 32], dtype=TensorProto.BOOL)


@pytest.mark.parametrize("op_name", ["BitwiseAnd", "BitwiseOr", "BitwiseXor"])
def test_bitwise(op_name: str):
    verify_binary(op_name, [32, 32], [32, 32], [32, 32], dtype=TensorProto.UINT64, opset=18)


def test_bitwise_not():
    verify_unary(
        "BitwiseNot",
        [32, 32],
        input_dtype=TensorProto.UINT64,
        output_dtype=TensorProto.UINT64,
        opset=18,
    )


@pytest.mark.parametrize("direction", ["LEFT", "RIGHT"])
def test_bitwise_shift(direction: str):
    shape = [32, 32]
    dtype = TensorProto.UINT64
    test_node = helper.make_node("BitShift", ["a", "b"], ["c"], direction=direction)
    graph = helper.make_graph(
        [test_node],
        "binary_test",
        inputs=[
            helper.make_tensor_value_info("a", dtype, shape),
            helper.make_tensor_value_info("b", dtype, shape),
        ],
        outputs=[helper.make_tensor_value_info("c", dtype, shape)],
    )

    model = helper.make_model(graph, producer_name="binary_test")
    check_correctness(model, inputs={"b": np.random.randint(0, 8, shape).astype("uint64")})


@pytest.mark.parametrize(
    "op_name",
    [
        "Sin",
        "Cos",
        "Tan",
        "Sinh",
        "Cosh",
        "Tanh",
        "Asin",
        "Acos",
        "Atan",
        "Asinh",
        "Acosh",
        "Atanh",
        "Neg",
        "Abs",
        "Log",
        "Exp",
        "Not",
        "Floor",
        "Ceil",
        "Round",
        "IsInf",
        "IsNaN",
        "Sqrt",
        "Relu",
        "Sign",
        "Softplus",
        "Erf",
        "Sigmoid",
        "Softmax",
        "LogSoftmax",
    ],
)
def test_unary(op_name: str):
    input_dtype = TensorProto.FLOAT
    if op_name in [
        "IsNaN",
        "IsInf",
    ]:
        pytest.skip(f"Skipping test {op_name} because current LegalizeOps does not support it.")
    elif op_name == "Not":
        input_dtype = TensorProto.BOOL
        output_dtype = TensorProto.BOOL
    else:
        output_dtype = TensorProto.FLOAT
    verify_unary(op_name, [8, 8, 8], input_dtype=input_dtype, output_dtype=output_dtype)


def test_reciprocal_ir():
    model = make_unary_model("Reciprocal", [2, 3])
    tvm_model = from_onnx(model, keep_params_in_input=True)

    @I.ir_module
    class Expected:
        @R.function
        def main(x: R.Tensor((2, 3), dtype="float32")) -> R.Tensor((2, 3), dtype="float32"):
            R.func_attr({"num_input": 1})
            with R.dataflow():
                gv: R.Tensor((2, 3), dtype="float32") = R.divide(R.const(1.0, "float32"), x)
                R.output(gv)
            return gv

    tvm.ir.assert_structural_equal(tvm_model, Expected)


def test_identity_ir():
    model = make_unary_model("Identity", [8, 8, 8])
    tvm_model = from_onnx(model, keep_params_in_input=True)

    @I.ir_module
    class Expected:
        @R.function
        def main(x: R.Tensor((8, 8, 8), dtype="float32")) -> R.Tensor((8, 8, 8), dtype="float32"):
            R.func_attr({"num_input": 1})
            with R.dataflow():
                gv: R.Tensor((8, 8, 8), dtype="float32") = x
                R.output(gv)
            return gv

    tvm.ir.assert_structural_equal(tvm_model, Expected)


def test_elu_ir():
    model = make_unary_model("Elu", [2, 3])
    tvm_model = from_onnx(model, keep_params_in_input=True)

    @I.ir_module
    class Expected:
        @R.function
        def main(x: R.Tensor((2, 3), dtype="float32")) -> R.Tensor((2, 3), dtype="float32"):
            R.func_attr({"num_input": 1})
            with R.dataflow():
                lv: R.Tensor((2, 3), dtype="float32") = R.exp(x)
                lv1: R.Tensor((2, 3), dtype="float32") = R.subtract(R.const(1.0, "float32"), lv)
                lv2: R.Tensor((2, 3), dtype="float32") = R.nn.relu(lv1)
                lv3: R.Tensor((2, 3), dtype="float32") = R.multiply(R.const(-1.0, "float32"), lv2)
                lv4: R.Tensor((2, 3), dtype="float32") = R.nn.relu(x)
                gv: R.Tensor((2, 3), dtype="float32") = R.add(lv3, lv4)
                R.output(gv)
            return gv

    tvm.ir.assert_structural_equal(tvm_model, Expected)


def test_hardswish_ir():
    model = make_unary_model("HardSwish", [2, 3])
    tvm_model = from_onnx(model, keep_params_in_input=True)

    @I.ir_module
    class Expected:
        @R.function
        def main(x: R.Tensor((2, 3), dtype="float32")) -> R.Tensor((2, 3), dtype="float32"):
            R.func_attr({"num_input": 1})
            with R.dataflow():
                lv: R.Tensor((2, 3), dtype="float32") = R.add(x, R.const(3.0, "float32"))
                lv1: R.Tensor((2, 3), dtype="float32") = R.clip(
                    lv, R.prim_value(0), R.prim_value(6)
                )
                lv2: R.Tensor((2, 3), dtype="float32") = R.divide(lv1, R.const(6.0, "float32"))
                gv: R.Tensor((2, 3), dtype="float32") = R.multiply(x, lv2)
                R.output(gv)
            return gv

    tvm.ir.assert_structural_equal(tvm_model, Expected)


def test_softsign_ir():
    model = make_unary_model("Softsign", [2, 3])
    tvm_model = from_onnx(model, keep_params_in_input=True)

    @I.ir_module
    class Expected:
        @R.function
        def main(x: R.Tensor((2, 3), dtype="float32")) -> R.Tensor((2, 3), dtype="float32"):
            R.func_attr({"num_input": 1})
            with R.dataflow():
                lv: R.Tensor((2, 3), dtype="float32") = R.abs(x)
                lv1: R.Tensor((2, 3), dtype="float32") = R.add(lv, R.const(1.0, "float32"))
                gv: R.Tensor((2, 3), dtype="float32") = R.divide(x, lv1)
                R.output(gv)
            return gv

    tvm.ir.assert_structural_equal(tvm_model, Expected)


def test_hardmax_ir():
    model = make_unary_model("Hardmax", [2, 3])
    tvm_model = from_onnx(model, keep_params_in_input=True)

    @I.ir_module
    class Expected:
        @R.function
        def main(x: R.Tensor((2, 3), dtype="float32")) -> R.Tensor((2, 3), dtype="float32"):
            R.func_attr({"num_input": 1})
            with R.dataflow():
                lv: R.Tensor((2,), dtype="int64") = R.argmax(x, axis=1, keepdims=False)
                gv: R.Tensor((2, 3), dtype="float32") = R.one_hot(
                    lv,
                    R.prim_value(T.float32(1.0)),
                    R.prim_value(T.float32(0.0)),
                    depth=3,
                    axis=1,
                )
                R.output(gv)
            return gv

    tvm.ir.assert_structural_equal(tvm_model, Expected)


def test_legacy_softmax_family_opset11_axis_semantics():
    def verify_legacy_softmax_family_axis_ir(op_name: str, expected, axis_attr: int | None = None):
        attrs = {} if axis_attr is None else {"axis": axis_attr}
        node = helper.make_node(op_name, ["x"], ["y"], **attrs)
        graph = helper.make_graph(
            [node],
            "legacy_softmax_family_axis_ir_test",
            inputs=[helper.make_tensor_value_info("x", TensorProto.FLOAT, [2, 3, 4])],
            outputs=[helper.make_tensor_value_info("y", TensorProto.FLOAT, [2, 3, 4])],
        )
        model = helper.make_model(
            graph,
            producer_name="legacy_softmax_family_axis_ir_test",
            opset_imports=[helper.make_opsetid("", 11)],
        )
        tvm_model = from_onnx(model, opset=11, keep_params_in_input=True)
        tvm.ir.assert_structural_equal(tvm_model, expected)

    @I.ir_module
    class ExpectedSoftmaxAxis0:
        @R.function
        def main(
            x: R.Tensor((2, 3, 4), dtype="float32"),
        ) -> R.Tensor((2, 3, 4), dtype="float32"):
            R.func_attr({"num_input": 1})
            with R.dataflow():
                lv: R.Tensor((1, 24), dtype="float32") = R.reshape(x, R.shape([1, 24]))
                lv1: R.Tensor((1, 24), dtype="float32") = R.nn.softmax(lv, axis=-1)
                gv: R.Tensor((2, 3, 4), dtype="float32") = R.reshape(lv1, R.shape([2, 3, 4]))
                R.output(gv)
            return gv

    @I.ir_module
    class ExpectedSoftmaxAxis1:
        @R.function
        def main(
            x: R.Tensor((2, 3, 4), dtype="float32"),
        ) -> R.Tensor((2, 3, 4), dtype="float32"):
            R.func_attr({"num_input": 1})
            with R.dataflow():
                lv: R.Tensor((2, 12), dtype="float32") = R.reshape(x, R.shape([2, 12]))
                lv1: R.Tensor((2, 12), dtype="float32") = R.nn.softmax(lv, axis=-1)
                gv: R.Tensor((2, 3, 4), dtype="float32") = R.reshape(lv1, R.shape([2, 3, 4]))
                R.output(gv)
            return gv

    @I.ir_module
    class ExpectedSoftmaxAxisRank:
        @R.function
        def main(
            x: R.Tensor((2, 3, 4), dtype="float32"),
        ) -> R.Tensor((2, 3, 4), dtype="float32"):
            R.func_attr({"num_input": 1})
            with R.dataflow():
                lv: R.Tensor((24, 1), dtype="float32") = R.reshape(x, R.shape([24, 1]))
                lv1: R.Tensor((24, 1), dtype="float32") = R.nn.softmax(lv, axis=-1)
                gv: R.Tensor((2, 3, 4), dtype="float32") = R.reshape(lv1, R.shape([2, 3, 4]))
                R.output(gv)
            return gv

    @I.ir_module
    class ExpectedLogSoftmaxAxis0:
        @R.function
        def main(
            x: R.Tensor((2, 3, 4), dtype="float32"),
        ) -> R.Tensor((2, 3, 4), dtype="float32"):
            R.func_attr({"num_input": 1})
            with R.dataflow():
                lv: R.Tensor((1, 24), dtype="float32") = R.reshape(x, R.shape([1, 24]))
                lv1: R.Tensor((1, 24), dtype="float32") = R.nn.log_softmax(lv, axis=-1)
                gv: R.Tensor((2, 3, 4), dtype="float32") = R.reshape(lv1, R.shape([2, 3, 4]))
                R.output(gv)
            return gv

    @I.ir_module
    class ExpectedLogSoftmaxAxis1:
        @R.function
        def main(
            x: R.Tensor((2, 3, 4), dtype="float32"),
        ) -> R.Tensor((2, 3, 4), dtype="float32"):
            R.func_attr({"num_input": 1})
            with R.dataflow():
                lv: R.Tensor((2, 12), dtype="float32") = R.reshape(x, R.shape([2, 12]))
                lv1: R.Tensor((2, 12), dtype="float32") = R.nn.log_softmax(lv, axis=-1)
                gv: R.Tensor((2, 3, 4), dtype="float32") = R.reshape(lv1, R.shape([2, 3, 4]))
                R.output(gv)
            return gv

    @I.ir_module
    class ExpectedLogSoftmaxAxisRank:
        @R.function
        def main(
            x: R.Tensor((2, 3, 4), dtype="float32"),
        ) -> R.Tensor((2, 3, 4), dtype="float32"):
            R.func_attr({"num_input": 1})
            with R.dataflow():
                lv: R.Tensor((24, 1), dtype="float32") = R.reshape(x, R.shape([24, 1]))
                lv1: R.Tensor((24, 1), dtype="float32") = R.nn.log_softmax(lv, axis=-1)
                gv: R.Tensor((2, 3, 4), dtype="float32") = R.reshape(lv1, R.shape([2, 3, 4]))
                R.output(gv)
            return gv

    @I.ir_module
    class ExpectedHardmaxAxis0:
        @R.function
        def main(
            x: R.Tensor((2, 3, 4), dtype="float32"),
        ) -> R.Tensor((2, 3, 4), dtype="float32"):
            R.func_attr({"num_input": 1})
            with R.dataflow():
                lv: R.Tensor((1, 24), dtype="float32") = R.reshape(x, R.shape([1, 24]))
                lv1: R.Tensor((1,), dtype="int64") = R.argmax(lv, axis=1, keepdims=False)
                lv2: R.Tensor((1, 24), dtype="float32") = R.one_hot(
                    lv1,
                    R.prim_value(T.float32(1.0)),
                    R.prim_value(T.float32(0.0)),
                    depth=24,
                    axis=1,
                )
                gv: R.Tensor((2, 3, 4), dtype="float32") = R.reshape(lv2, R.shape([2, 3, 4]))
                R.output(gv)
            return gv

    @I.ir_module
    class ExpectedHardmaxAxis1:
        @R.function
        def main(
            x: R.Tensor((2, 3, 4), dtype="float32"),
        ) -> R.Tensor((2, 3, 4), dtype="float32"):
            R.func_attr({"num_input": 1})
            with R.dataflow():
                lv: R.Tensor((2, 12), dtype="float32") = R.reshape(x, R.shape([2, 12]))
                lv1: R.Tensor((2,), dtype="int64") = R.argmax(lv, axis=1, keepdims=False)
                lv2: R.Tensor((2, 12), dtype="float32") = R.one_hot(
                    lv1,
                    R.prim_value(T.float32(1.0)),
                    R.prim_value(T.float32(0.0)),
                    depth=12,
                    axis=1,
                )
                gv: R.Tensor((2, 3, 4), dtype="float32") = R.reshape(lv2, R.shape([2, 3, 4]))
                R.output(gv)
            return gv

    @I.ir_module
    class ExpectedHardmaxAxisRank:
        @R.function
        def main(
            x: R.Tensor((2, 3, 4), dtype="float32"),
        ) -> R.Tensor((2, 3, 4), dtype="float32"):
            R.func_attr({"num_input": 1})
            with R.dataflow():
                lv: R.Tensor((24, 1), dtype="float32") = R.reshape(x, R.shape([24, 1]))
                lv1: R.Tensor((24,), dtype="int64") = R.argmax(lv, axis=1, keepdims=False)
                lv2: R.Tensor((24, 1), dtype="float32") = R.one_hot(
                    lv1,
                    R.prim_value(T.float32(1.0)),
                    R.prim_value(T.float32(0.0)),
                    depth=1,
                    axis=1,
                )
                gv: R.Tensor((2, 3, 4), dtype="float32") = R.reshape(lv2, R.shape([2, 3, 4]))
                R.output(gv)
            return gv

    # Default axis and equivalent negative axis both flatten from axis 1.
    verify_legacy_softmax_family_axis_ir("Softmax", ExpectedSoftmaxAxis1)
    verify_legacy_softmax_family_axis_ir("LogSoftmax", ExpectedLogSoftmaxAxis1)
    verify_legacy_softmax_family_axis_ir("Hardmax", ExpectedHardmaxAxis1)
    verify_legacy_softmax_family_axis_ir("Softmax", ExpectedSoftmaxAxis1, axis_attr=-2)
    verify_legacy_softmax_family_axis_ir("LogSoftmax", ExpectedLogSoftmaxAxis1, axis_attr=-2)
    verify_legacy_softmax_family_axis_ir("Hardmax", ExpectedHardmaxAxis1, axis_attr=-2)

    # Positive axis 0 flattens the whole input as one row.
    verify_legacy_softmax_family_axis_ir("Softmax", ExpectedSoftmaxAxis0, axis_attr=0)
    verify_legacy_softmax_family_axis_ir("LogSoftmax", ExpectedLogSoftmaxAxis0, axis_attr=0)
    verify_legacy_softmax_family_axis_ir("Hardmax", ExpectedHardmaxAxis0, axis_attr=0)

    # Axis equal to rank produces a trailing singleton reduction dimension.
    verify_legacy_softmax_family_axis_ir("Softmax", ExpectedSoftmaxAxisRank, axis_attr=3)
    verify_legacy_softmax_family_axis_ir("LogSoftmax", ExpectedLogSoftmaxAxisRank, axis_attr=3)
    verify_legacy_softmax_family_axis_ir("Hardmax", ExpectedHardmaxAxisRank, axis_attr=3)


@pytest.mark.parametrize("op_name", ["Softmax", "LogSoftmax"])
def test_softmax_family_opset13_default_axis_semantics(op_name: str):
    verify_unary(op_name, [2, 3, 4], opset=13)


def test_hardmax_opset13_default_axis_ir():
    model = make_unary_model("Hardmax", [2, 3, 4])
    model.opset_import[0].version = 13
    tvm_model = from_onnx(model, opset=13, keep_params_in_input=True)

    @I.ir_module
    class Expected:
        @R.function
        def main(
            x: R.Tensor((2, 3, 4), dtype="float32"),
        ) -> R.Tensor((2, 3, 4), dtype="float32"):
            R.func_attr({"num_input": 1})
            with R.dataflow():
                lv: R.Tensor((2, 3), dtype="int64") = R.argmax(x, axis=2, keepdims=False)
                gv: R.Tensor((2, 3, 4), dtype="float32") = R.one_hot(
                    lv,
                    R.prim_value(T.float32(1.0)),
                    R.prim_value(T.float32(0.0)),
                    depth=4,
                    axis=2,
                )
                R.output(gv)
            return gv

    tvm.ir.assert_structural_equal(tvm_model, Expected)


def test_legacy_softmax_family_opset1_ir_semantics():
    def verify_legacy_softmax_family_opset1_ir(op_name: str, expected):
        node = helper.make_node(op_name, ["x"], ["y"])
        graph = helper.make_graph(
            [node],
            "legacy_softmax_family_opset1_ir_test",
            inputs=[helper.make_tensor_value_info("x", TensorProto.FLOAT, [2, 3, 4])],
            outputs=[helper.make_tensor_value_info("y", TensorProto.FLOAT, [2, 3, 4])],
        )
        model = helper.make_model(
            graph,
            producer_name="legacy_softmax_family_opset1_ir_test",
            opset_imports=[helper.make_opsetid("", 1)],
        )
        tvm_model = from_onnx(model, opset=1, keep_params_in_input=True)
        tvm.ir.assert_structural_equal(tvm_model, expected)

    @I.ir_module
    class ExpectedSoftmax:
        @R.function
        def main(
            x: R.Tensor((2, 3, 4), dtype="float32"),
        ) -> R.Tensor((2, 3, 4), dtype="float32"):
            R.func_attr({"num_input": 1})
            with R.dataflow():
                lv: R.Tensor((2, 12), dtype="float32") = R.reshape(x, R.shape([2, 12]))
                lv1: R.Tensor((2, 12), dtype="float32") = R.nn.softmax(lv, axis=-1)
                gv: R.Tensor((2, 3, 4), dtype="float32") = R.reshape(lv1, R.shape([2, 3, 4]))
                R.output(gv)
            return gv

    @I.ir_module
    class ExpectedLogSoftmax:
        @R.function
        def main(
            x: R.Tensor((2, 3, 4), dtype="float32"),
        ) -> R.Tensor((2, 3, 4), dtype="float32"):
            R.func_attr({"num_input": 1})
            with R.dataflow():
                lv: R.Tensor((2, 12), dtype="float32") = R.reshape(x, R.shape([2, 12]))
                lv1: R.Tensor((2, 12), dtype="float32") = R.nn.log_softmax(lv, axis=-1)
                gv: R.Tensor((2, 3, 4), dtype="float32") = R.reshape(lv1, R.shape([2, 3, 4]))
                R.output(gv)
            return gv

    @I.ir_module
    class ExpectedHardmax:
        @R.function
        def main(
            x: R.Tensor((2, 3, 4), dtype="float32"),
        ) -> R.Tensor((2, 3, 4), dtype="float32"):
            R.func_attr({"num_input": 1})
            with R.dataflow():
                lv: R.Tensor((2, 12), dtype="float32") = R.reshape(x, R.shape([2, 12]))
                lv1: R.Tensor((2,), dtype="int64") = R.argmax(lv, axis=1, keepdims=False)
                lv2: R.Tensor((2, 12), dtype="float32") = R.one_hot(
                    lv1,
                    R.prim_value(T.float32(1.0)),
                    R.prim_value(T.float32(0.0)),
                    depth=12,
                    axis=1,
                )
                gv: R.Tensor((2, 3, 4), dtype="float32") = R.reshape(lv2, R.shape([2, 3, 4]))
                R.output(gv)
            return gv

    verify_legacy_softmax_family_opset1_ir("Softmax", ExpectedSoftmax)
    verify_legacy_softmax_family_opset1_ir("LogSoftmax", ExpectedLogSoftmax)
    verify_legacy_softmax_family_opset1_ir("Hardmax", ExpectedHardmax)


def test_round_ties_to_even():
    """ONNX Round must use ties-to-even (banker's rounding), not ties-away-from-zero.

    Per the ONNX spec: "For cases where number is exactly halfway between two
    integers, it rounds to the nearest even integer."
    https://onnx.ai/onnx/operators/onnx__Round.html
    """
    round_node = helper.make_node("Round", ["x"], ["y"])
    graph = helper.make_graph(
        [round_node],
        "round_ties_to_even_test",
        inputs=[helper.make_tensor_value_info("x", TensorProto.FLOAT, [6])],
        outputs=[helper.make_tensor_value_info("y", TensorProto.FLOAT, [6])],
    )
    model = helper.make_model(graph, producer_name="round_ties_to_even_test")
    # Midpoint values: 0.5->0, 1.5->2, 2.5->2, -0.5->0, -1.5->-2, -2.5->-2 (ties-to-even)
    # Ties-away would give: 0.5->1, 1.5->2, 2.5->3, -0.5->-1, -1.5->-2, -2.5->-3
    inputs = {"x": np.array([0.5, 1.5, 2.5, -0.5, -1.5, -2.5], dtype="float32")}
    check_correctness(model, inputs=inputs, opset=11)


@pytest.mark.parametrize("from_type", [TensorProto.INT32, TensorProto.FLOAT, TensorProto.FLOAT16])
@pytest.mark.parametrize("to_type", [TensorProto.INT32, TensorProto.FLOAT, TensorProto.FLOAT16])
def test_cast(from_type, to_type):
    cast_node = helper.make_node("Cast", ["a"], ["a_float"], to=to_type)

    graph = helper.make_graph(
        [cast_node],
        "cast_test",
        inputs=[
            helper.make_tensor_value_info("a", from_type, [1, 32]),
        ],
        outputs=[helper.make_tensor_value_info("a_float", to_type, [1, 32])],
    )

    model = helper.make_model(graph, producer_name="cast_test")
    check_correctness(model, opset=13)


@pytest.mark.parametrize("to_type", [TensorProto.INT64, TensorProto.UINT64])
def test_cast_float_to_64bit_int_dynamic(to_type):
    cast_node = helper.make_node("Cast", ["a"], ["b"], to=to_type)
    graph = helper.make_graph(
        [cast_node],
        "cast_float_to_64bit_int_dynamic_test",
        inputs=[helper.make_tensor_value_info("a", TensorProto.FLOAT, [1, 8])],
        outputs=[helper.make_tensor_value_info("b", to_type, [1, 8])],
    )
    model = helper.make_model(graph, producer_name="cast_float_to_64bit_int_dynamic_test")
    inputs = {"a": np.array([[0.0, 1.2, 2.8, 7.9, 15.1, 31.7, 63.4, 127.9]], dtype=np.float32)}
    check_correctness(model, inputs=inputs, opset=13, check_dtypes=True)


def test_cast_nan_inf_to_int8():
    vals = np.array([300.0, np.nan, np.inf, -np.inf, 50.0, -50.0], dtype=np.float32)
    node = helper.make_node("Cast", inputs=["a"], outputs=["b"], to=TensorProto.INT8)
    graph = helper.make_graph(
        [node],
        "cast_nan_inf_test",
        inputs=[helper.make_tensor_value_info("a", TensorProto.FLOAT, list(vals.shape))],
        outputs=[helper.make_tensor_value_info("b", TensorProto.INT8, list(vals.shape))],
    )
    model = helper.make_model(graph, producer_name="cast_nan_inf_test")
    tvm_output = run_in_tvm(model, inputs={"a": vals}, opset=13)
    out_np = tvm_output.numpy()
    expected = np.array([44, 0, 0, 0, 50, -50], dtype=np.int8)
    assert out_np.dtype == np.int8
    np.testing.assert_array_equal(out_np, expected)


def test_gather():
    def _verify_gather(data_shape, indices, out_shape, expected, axis=0):
        gather_node = helper.make_node("Gather", ["data", "indices"], ["y"], axis=axis)

        if isinstance(indices, list | tuple):
            indices_shape = np.asarray(indices).shape
        else:
            indices_shape = []

        graph = helper.make_graph(
            [gather_node],
            "gather_test",
            inputs=[
                helper.make_tensor_value_info("data", TensorProto.FLOAT, data_shape),
                helper.make_tensor_value_info("indices", TensorProto.INT64, indices_shape),
            ],
            outputs=[helper.make_tensor_value_info("y", TensorProto.FLOAT, out_shape)],
        )

        model = helper.make_model(
            graph, producer_name="gather_test", opset_imports=[helper.make_opsetid("", 14)]
        )
        tvm_model = from_onnx(model, opset=14, keep_params_in_input=True)
        tvm.ir.assert_structural_equal(tvm_model, expected)

    @I.ir_module
    class ExpectedRank4Axis0:
        @R.function
        def main(
            data: R.Tensor((5, 4, 3, 2), dtype="float32"),
            indices: R.Tensor((3,), dtype="int64"),
        ) -> R.Tensor((3, 4, 3, 2), dtype="float32"):
            R.func_attr({"num_input": 2})
            with R.dataflow():
                lv: R.Shape([5, 4, 3, 2]) = R.shape_of(data)
                lv1: R.Tensor((4,), dtype="int64") = R.shape_to_tensor(lv)
                lv2: R.Tensor((3,), dtype="bool") = R.less(indices, R.const(0, "int64"))
                lv3: R.Tensor((), dtype="int64") = R.take(
                    lv1, R.const(0, "int64"), axis=0, mode="wrap"
                )
                lv4: R.Tensor((3,), dtype="int64") = R.add(indices, lv3)
                lv5: R.Tensor((3,), dtype="int64") = R.where(lv2, lv4, indices)
                gv: R.Tensor((3, 4, 3, 2), dtype="float32") = R.take(data, lv5, axis=0, mode="fast")
                R.output(gv)
            return gv

    @I.ir_module
    class ExpectedScalarIndex:
        @R.function
        def main(
            data: R.Tensor((3,), dtype="float32"),
            indices: R.Tensor((), dtype="int64"),
        ) -> R.Tensor((), dtype="float32"):
            R.func_attr({"num_input": 2})
            with R.dataflow():
                lv: R.Shape([3]) = R.shape_of(data)
                lv1: R.Tensor((1,), dtype="int64") = R.shape_to_tensor(lv)
                lv2: R.Tensor((), dtype="bool") = R.less(indices, R.const(0, "int64"))
                lv3: R.Tensor((), dtype="int64") = R.take(
                    lv1, R.const(0, "int64"), axis=0, mode="wrap"
                )
                lv4: R.Tensor((), dtype="int64") = R.add(indices, lv3)
                lv5: R.Tensor((), dtype="int64") = R.where(lv2, lv4, indices)
                gv: R.Tensor((), dtype="float32") = R.take(data, lv5, axis=0, mode="fast")
                R.output(gv)
            return gv

    @I.ir_module
    class ExpectedRank2Axis1:
        @R.function
        def main(
            data: R.Tensor((3, 3), dtype="float32"),
            indices: R.Tensor((1, 2), dtype="int64"),
        ) -> R.Tensor((3, 1, 2), dtype="float32"):
            R.func_attr({"num_input": 2})
            with R.dataflow():
                lv: R.Shape([3, 3]) = R.shape_of(data)
                lv1: R.Tensor((2,), dtype="int64") = R.shape_to_tensor(lv)
                lv2: R.Tensor((1, 2), dtype="bool") = R.less(indices, R.const(0, "int64"))
                lv3: R.Tensor((), dtype="int64") = R.take(
                    lv1, R.const(1, "int64"), axis=0, mode="wrap"
                )
                lv4: R.Tensor((1, 2), dtype="int64") = R.add(indices, lv3)
                lv5: R.Tensor((1, 2), dtype="int64") = R.where(lv2, lv4, indices)
                gv: R.Tensor((3, 1, 2), dtype="float32") = R.take(data, lv5, axis=1, mode="fast")
                R.output(gv)
            return gv

    _verify_gather([5, 4, 3, 2], [0, 1, 3], [3, 4, 3, 2], ExpectedRank4Axis0)
    _verify_gather([3], 0, [], ExpectedScalarIndex)
    _verify_gather([3, 3], [[0, 2]], [3, 1, 2], ExpectedRank2Axis1, 1)


def _make_gather_negative_indices_expected(axis: int, indices_shape, indices_type):
    indices_shape = tuple(indices_shape)
    indices_dtype = "int64" if indices_type == TensorProto.INT64 else "int32"

    if indices_type == TensorProto.INT64:

        @I.ir_module
        class ExpectedGatherNegativeInt64:
            @R.function
            def main(
                data: R.Tensor((3, 4), dtype="float32"),
                indices: R.Tensor(indices_shape, dtype=indices_dtype),
            ):
                R.func_attr({"num_input": 2})
                with R.dataflow():
                    lv: R.Shape([3, 4]) = R.shape_of(data)
                    lv1: R.Tensor((2,), dtype="int64") = R.shape_to_tensor(lv)
                    lv2 = R.less(indices, R.const(0, "int64"))
                    lv3: R.Tensor((), dtype="int64") = R.take(
                        lv1, R.const(axis, "int64"), axis=0, mode="wrap"
                    )
                    lv4 = R.add(indices, lv3)
                    lv5 = R.where(lv2, lv4, indices)
                    gv = R.take(data, lv5, axis=axis, mode="fast")
                    R.output(gv)
                return gv

        return ExpectedGatherNegativeInt64

    if indices_type == TensorProto.INT32:

        @I.ir_module
        class ExpectedGatherNegativeInt32:
            @R.function
            def main(
                data: R.Tensor((3, 4), dtype="float32"),
                indices: R.Tensor(indices_shape, dtype=indices_dtype),
            ):
                R.func_attr({"num_input": 2})
                with R.dataflow():
                    lv: R.Shape([3, 4]) = R.shape_of(data)
                    lv1: R.Tensor((2,), dtype="int64") = R.shape_to_tensor(lv)
                    lv2: R.Tensor((), dtype="int64") = R.take(
                        lv1, R.const(axis, "int64"), axis=0, mode="wrap"
                    )
                    lv3 = R.less(indices, R.const(0, "int32"))
                    lv4: R.Tensor((), dtype="int32") = R.astype(lv2, dtype="int32")
                    lv5 = R.add(indices, lv4)
                    lv6 = R.where(lv3, lv5, indices)
                    gv = R.take(data, lv6, axis=axis, mode="fast")
                    R.output(gv)
                return gv

        return ExpectedGatherNegativeInt32

    raise AssertionError(
        f"Unexpected Gather negative-index case: axis={axis}, "
        f"indices_shape={indices_shape}, indices_type={indices_type}"
    )


def test_gather_negative_indices():
    def verify_gather_negative_indices(axis, indices, out_shape, indices_type, expected):
        gather_node = helper.make_node("Gather", ["data", "indices"], ["y"], axis=axis)
        indices_shape = np.asarray(indices).shape

        graph = helper.make_graph(
            [gather_node],
            "gather_negative_indices_test",
            inputs=[
                helper.make_tensor_value_info("data", TensorProto.FLOAT, [3, 4]),
                helper.make_tensor_value_info("indices", indices_type, indices_shape),
            ],
            outputs=[helper.make_tensor_value_info("y", TensorProto.FLOAT, out_shape)],
        )

        model = helper.make_model(
            graph,
            producer_name="gather_negative_indices_test",
            opset_imports=[helper.make_opsetid("", 14)],
        )
        tvm_model = from_onnx(model, opset=14, keep_params_in_input=True)
        tvm.ir.assert_structural_equal(tvm_model, expected)

    for axis, indices, out_shape, indices_type in [
        (0, [-1, 0], [2, 4], TensorProto.INT64),
        (1, [-1, 0], [3, 2], TensorProto.INT64),
        (1, [[-1, 0], [1, -2]], [3, 2, 2], TensorProto.INT64),
        (0, [-1, 0], [2, 4], TensorProto.INT32),
        (1, [-1, 0], [3, 2], TensorProto.INT32),
        (1, [[-1, 0], [1, -2]], [3, 2, 2], TensorProto.INT32),
    ]:
        verify_gather_negative_indices(
            axis,
            indices,
            out_shape,
            indices_type,
            _make_gather_negative_indices_expected(axis, np.asarray(indices).shape, indices_type),
        )


def test_gather_negative_indices_ir_normalization():
    def verify_gather_negative_indices_ir_normalization(indices_type, expected):
        gather_node = helper.make_node("Gather", ["data", "indices"], ["y"], axis=1)
        graph = helper.make_graph(
            [gather_node],
            "gather_negative_indices_ir_test",
            inputs=[
                helper.make_tensor_value_info("data", TensorProto.FLOAT, [3, 4]),
                helper.make_tensor_value_info("indices", indices_type, [2]),
            ],
            outputs=[helper.make_tensor_value_info("y", TensorProto.FLOAT, [3, 2])],
        )

        model = helper.make_model(graph, producer_name="gather_negative_indices_ir_test")
        tvm_model = from_onnx(model, opset=13, keep_params_in_input=True)
        tvm.ir.assert_structural_equal(tvm_model, expected)

    verify_gather_negative_indices_ir_normalization(
        TensorProto.INT64, _make_gather_negative_indices_expected(1, (2,), TensorProto.INT64)
    )
    verify_gather_negative_indices_ir_normalization(
        TensorProto.INT32, _make_gather_negative_indices_expected(1, (2,), TensorProto.INT32)
    )


@pytest.mark.parametrize(
    "data_shape, indices_shape, axis",
    [
        ([3, 4, 5], [1, 4, 5], 0),
        ([3, 4, 5], [3, 2, 5], 1),
        ([3, 4, 5], [3, 4, 2], 2),
    ],
)
def test_gather_elements(data_shape, indices_shape, axis):
    gather_elements_node = helper.make_node("GatherElements", ["data", "indices"], ["y"], axis=axis)

    graph = helper.make_graph(
        [gather_elements_node],
        "gather_elements_test",
        inputs=[
            helper.make_tensor_value_info("data", TensorProto.FLOAT, data_shape),
            helper.make_tensor_value_info("indices", TensorProto.INT64, indices_shape),
        ],
        outputs=[helper.make_tensor_value_info("y", TensorProto.FLOAT, indices_shape)],
    )

    model = helper.make_model(graph, producer_name="gather_elements_test")
    input_values = {
        "data": np.random.randn(*data_shape).astype("float32"),
        "indices": np.random.randint(0, data_shape[axis], indices_shape).astype("int64"),
    }
    check_correctness(model, inputs=input_values)


@pytest.mark.parametrize(
    "data_shape, indices_shape, batch_dims",
    [
        ([2, 2], [2, 2], 0),
        ([2, 2], [2, 1], 0),
        ([2, 2, 2], [1], 0),
        ([2, 2, 2], [2, 2], 0),
        ([2, 2, 2], [2, 1, 2], 0),
        ([2, 2, 2], [2, 2], 1),
        ([2, 2, 2], [2, 1], 1),
    ],
)
def test_gather_nd(data_shape, indices_shape, batch_dims):
    gather_nd_node = helper.make_node("GatherND", ["data", "indices"], ["y"], batch_dims=batch_dims)

    graph = helper.make_graph(
        [gather_nd_node],
        "gather_nd_test",
        inputs=[
            helper.make_tensor_value_info("data", TensorProto.FLOAT, data_shape),
            helper.make_tensor_value_info("indices", TensorProto.INT64, indices_shape),
        ],
        outputs=[helper.make_tensor_value_info("y", TensorProto.FLOAT, None)],
    )

    model = helper.make_model(graph, producer_name="gather_nd_test")
    input_values = {
        "data": np.random.randn(*data_shape).astype("float32"),
        "indices": np.random.randint(0, 2, indices_shape).astype("int64"),
    }
    check_correctness(model, inputs=input_values)


@pytest.mark.parametrize("axis", [0, 1, 2])
@pytest.mark.parametrize(("name", "opset"), [("Scatter", 10), ("ScatterElements", 11)])
def test_scatter(axis: int, name: str, opset: int):
    if axis != 1:
        pytest.skip("The current topi impl is wrong, which only works for axis=1")
    input_shape = [16, 16, 16]
    indices_shape = [8, 8, 8]
    updates_shape = [8, 8, 8]
    output_shape = [16, 16, 16]
    node = helper.make_node(name, ["data", "indices", "updates"], ["output"], axis=axis)
    graph = helper.make_graph(
        [node],
        "scatter_test",
        inputs=[
            helper.make_tensor_value_info("data", TensorProto.FLOAT, input_shape),
            helper.make_tensor_value_info("indices", TensorProto.INT64, indices_shape),
            helper.make_tensor_value_info("updates", TensorProto.FLOAT, updates_shape),
        ],
        outputs=[helper.make_tensor_value_info("output", TensorProto.FLOAT, output_shape)],
    )
    model = helper.make_model(graph, producer_name="scatter_test")
    indices = np.random.randint(0, 16, indices_shape)
    check_correctness(model, inputs={"indices": indices}, opset=opset)


@pytest.mark.parametrize(
    "reduction, opset, data, indices, updates",
    [
        (
            None,
            11,
            np.array([[1, 2, 3], [4, 5, 6]], dtype="float32"),
            np.array([[2, 0, 1], [1, 2, 0]], dtype="int64"),
            np.array([[30, 10, 20], [50, 60, 40]], dtype="float32"),
        ),
        (
            "none",
            18,
            np.array([[1, 2, 3], [4, 5, 6]], dtype="float32"),
            np.array([[2, 0, 1], [1, 2, 0]], dtype="int64"),
            np.array([[30, 10, 20], [50, 60, 40]], dtype="float32"),
        ),
        (
            "add",
            16,
            np.full((2, 3), 10, dtype="float32"),
            np.array([[0, 0, 2], [1, 1, 2]], dtype="int64"),
            np.array([[2, 5, 7], [20, 3, 4]], dtype="float32"),
        ),
        (
            "mul",
            16,
            np.full((2, 3), 10, dtype="float32"),
            np.array([[0, 0, 2], [1, 1, 2]], dtype="int64"),
            np.array([[2, 5, 7], [20, 3, 4]], dtype="float32"),
        ),
        (
            "min",
            18,
            np.full((2, 3), 10, dtype="float32"),
            np.array([[0, 0, 2], [1, 1, 2]], dtype="int64"),
            np.array([[2, 5, 7], [20, 3, 4]], dtype="float32"),
        ),
        (
            "max",
            18,
            np.full((2, 3), 10, dtype="float32"),
            np.array([[0, 0, 2], [1, 1, 2]], dtype="int64"),
            np.array([[2, 5, 7], [20, 3, 4]], dtype="float32"),
        ),
    ],
)
def test_scatter_elements_reduction(reduction, opset, data, indices, updates):
    attrs = {"axis": 1}
    if reduction is not None:
        attrs["reduction"] = reduction
    scatter_elements_node = helper.make_node(
        "ScatterElements", ["data", "indices", "updates"], ["output"], **attrs
    )

    graph = helper.make_graph(
        [scatter_elements_node],
        "scatter_elements_reduction_test",
        inputs=[
            helper.make_tensor_value_info("data", TensorProto.FLOAT, list(data.shape)),
            helper.make_tensor_value_info("indices", TensorProto.INT64, list(indices.shape)),
            helper.make_tensor_value_info("updates", TensorProto.FLOAT, list(updates.shape)),
        ],
        outputs=[helper.make_tensor_value_info("output", TensorProto.FLOAT, list(data.shape))],
    )
    model = helper.make_model(graph, producer_name="scatter_elements_reduction_test")

    check_correctness(
        model,
        inputs={"data": data, "indices": indices, "updates": updates},
        opset=opset,
    )


def test_scatter_elements_invalid_reduction():
    data_shape = [2, 3]
    scatter_elements_node = helper.make_node(
        "ScatterElements",
        ["data", "indices", "updates"],
        ["output"],
        axis=1,
        reduction="unsupported",
    )

    graph = helper.make_graph(
        [scatter_elements_node],
        "scatter_elements_invalid_reduction_test",
        inputs=[
            helper.make_tensor_value_info("data", TensorProto.FLOAT, data_shape),
            helper.make_tensor_value_info("indices", TensorProto.INT64, data_shape),
            helper.make_tensor_value_info("updates", TensorProto.FLOAT, data_shape),
        ],
        outputs=[helper.make_tensor_value_info("output", TensorProto.FLOAT, data_shape)],
    )
    model = helper.make_model(graph, producer_name="scatter_elements_invalid_reduction_test")

    with pytest.raises(ValueError, match="Only .* reductions are supported, but got unsupported"):
        from_onnx(model, opset=18, keep_params_in_input=True)


@pytest.mark.parametrize("reduction", ["none", "add", "mul"])
def test_scatter_nd(reduction):
    def verify_scatter_nd(data_shape, indices_shape, updates_shape):
        scatter_nd_node = helper.make_node(
            "ScatterND",
            ["data", "indices", "updates"],
            ["output"],
            reduction=reduction,
        )

        graph = helper.make_graph(
            [scatter_nd_node],
            "scatter_nd_test",
            inputs=[
                helper.make_tensor_value_info("data", TensorProto.FLOAT, data_shape),
                helper.make_tensor_value_info("indices", TensorProto.INT64, indices_shape),
                helper.make_tensor_value_info("updates", TensorProto.FLOAT, updates_shape),
            ],
            outputs=[helper.make_tensor_value_info("output", TensorProto.FLOAT, data_shape)],
        )

        model = helper.make_model(graph, producer_name="scatter_nd_test")

        indices = np.random.choice(data_shape[0], indices_shape)
        check_correctness(model, inputs={"indices": indices}, opset=16)

    verify_scatter_nd([8], [4, 1], [4])
    verify_scatter_nd([4, 4, 4], [2, 1], [2, 4, 4])
    verify_scatter_nd([4, 5, 6], [2, 3, 2], [2, 3, 6])
    verify_scatter_nd([10], [5, 1], [5])


def test_compress():
    def verify_compress(
        tensor_shape: list[int],
        condition_shape: list[int] | None,
        axis: int | None,
        expected,
    ):
        if condition_shape is None:
            condition_shape = [tensor_shape[axis]]
        compress_node = helper.make_node("Compress", ["tensor", "condition"], ["output"], axis=axis)
        graph = helper.make_graph(
            [compress_node],
            "compress_test",
            inputs=[
                helper.make_tensor_value_info("tensor", TensorProto.FLOAT, tensor_shape),
                helper.make_tensor_value_info("condition", TensorProto.BOOL, condition_shape),
            ],
            outputs=[
                helper.make_tensor_value_info("output", TensorProto.FLOAT, [])
            ],  # shape is unknown
        )
        model = helper.make_model(graph, producer_name="compress_test")
        tvm_model = from_onnx(model, opset=11, keep_params_in_input=True)
        tvm.ir.assert_structural_equal(tvm_model, expected)

    def make_expected(tensor_shape: list[int], condition_shape: list[int] | None, axis: int | None):
        if condition_shape is None:
            condition_shape = [tensor_shape[axis]]
        tensor_shape = tuple(tensor_shape)
        condition_shape = tuple(condition_shape)

        if axis is None:
            flat_shape = (int(np.prod(tensor_shape)),)

            @I.ir_module
            class ExpectedCompressFlat:
                @R.function
                def main(
                    tensor: R.Tensor(tensor_shape, dtype="float32"),
                    condition: R.Tensor(condition_shape, dtype="bool"),
                ):
                    num_nonzero = T.int64()
                    R.func_attr({"num_input": 2})
                    with R.dataflow():
                        lv: R.Tensor((1, num_nonzero), dtype="int64") = R.match_cast(
                            R.nonzero(condition), R.Tensor((1, num_nonzero), dtype="int64")
                        )
                        lv1 = R.reshape(tensor, R.shape(flat_shape))
                        lv2: R.Tensor((num_nonzero,), dtype="int64") = R.reshape(
                            lv, R.shape([num_nonzero])
                        )
                        gv = R.take(lv1, lv2, axis=0, mode="fast")
                        R.output(gv)
                    return gv

            return ExpectedCompressFlat

        @I.ir_module
        class ExpectedCompressAxis:
            @R.function
            def main(
                tensor: R.Tensor(tensor_shape, dtype="float32"),
                condition: R.Tensor(condition_shape, dtype="bool"),
            ):
                num_nonzero = T.int64()
                R.func_attr({"num_input": 2})
                with R.dataflow():
                    lv: R.Tensor((1, num_nonzero), dtype="int64") = R.match_cast(
                        R.nonzero(condition), R.Tensor((1, num_nonzero), dtype="int64")
                    )
                    lv1: R.Tensor((num_nonzero,), dtype="int64") = R.reshape(
                        lv, R.shape([num_nonzero])
                    )
                    gv = R.take(tensor, lv1, axis=axis, mode="fast")
                    R.output(gv)
                return gv

        return ExpectedCompressAxis

    for tensor_shape, condition_shape, axis in [
        ([32, 32], [8], None),
        ([32, 32], [16], None),
        ([32, 32], [8], 0),
        ([32, 32], [16], 0),
        ([32, 32], None, 0),
        ([32, 32], [8], 1),
        ([32, 32], [16], 1),
        ([32, 32], None, 1),
    ]:
        verify_compress(
            tensor_shape, condition_shape, axis, make_expected(tensor_shape, condition_shape, axis)
        )


def test_size():
    test_node = helper.make_node("Size", ["x"], ["y"])
    input_shape = [3, 3, 3]
    graph = helper.make_graph(
        [test_node],
        "size_test",
        inputs=[helper.make_tensor_value_info("x", TensorProto.FLOAT, input_shape)],
        outputs=[helper.make_tensor_value_info("y", TensorProto.INT64, [3])],
    )

    model = helper.make_model(graph, producer_name="size_test")
    tvm_model = from_onnx(model, keep_params_in_input=True)

    @I.ir_module
    class Expected:
        @R.function
        def main(x: R.Tensor((3, 3, 3), dtype="float32")) -> R.Tensor((), dtype="int64"):
            R.func_attr({"num_input": 1})
            with R.dataflow():
                gv: R.Tensor((), dtype="int64") = R.size(x)
                R.output(gv)
            return gv

    tvm.ir.assert_structural_equal(tvm_model, Expected)


@pytest.mark.parametrize("k", [-1, 0, 1])
def test_eye_like(k: int):
    node = helper.make_node("EyeLike", ["x"], ["y"], k=k)
    graph = helper.make_graph(
        [node],
        "eye_like_test",
        inputs=[helper.make_tensor_value_info("x", TensorProto.FLOAT, [32, 32])],
        outputs=[helper.make_tensor_value_info("y", TensorProto.FLOAT, [32, 32])],
    )
    model = helper.make_model(graph, producer_name="eye_like_test")
    check_correctness(model)


def test_gemm():
    def verify_gemm(alpha, beta, useC, expected):
        if useC:
            gemm_node = helper.make_node(
                "Gemm", ["a", "b", "c"], ["y"], alpha=alpha, beta=beta, transA=1, transB=1
            )
        else:
            gemm_node = helper.make_node(
                "Gemm", ["a", "b"], ["y"], alpha=alpha, beta=beta, transA=1, transB=1
            )

        inputs = [
            helper.make_tensor_value_info("a", TensorProto.FLOAT, [4, 3]),
            helper.make_tensor_value_info("b", TensorProto.FLOAT, [5, 4]),
        ]
        if useC:
            inputs.append(helper.make_tensor_value_info("c", TensorProto.FLOAT, [1, 5]))

        graph = helper.make_graph(
            [gemm_node],
            "gemm_test",
            inputs=inputs,
            outputs=[helper.make_tensor_value_info("y", TensorProto.FLOAT, [3, 5])],
        )

        model = helper.make_model(graph, producer_name="gemm_test")
        tvm_model = from_onnx(model, keep_params_in_input=True)
        tvm.ir.assert_structural_equal(tvm_model, expected)

    def make_expected(alpha, beta, useC):
        alpha = 1.0 if alpha is None else alpha
        beta = 1.0 if beta is None else beta
        alpha = float(np.float32(alpha))
        beta = float(np.float32(beta))

        if not useC and alpha != 1.0:

            @I.ir_module
            class ExpectedScaledA:
                @R.function
                def main(
                    a: R.Tensor((4, 3), dtype="float32"),
                    b: R.Tensor((5, 4), dtype="float32"),
                ) -> R.Tensor((3, 5), dtype="float32"):
                    R.func_attr({"num_input": 2})
                    with R.dataflow():
                        lv: R.Tensor((4, 3), dtype="float32") = R.multiply(
                            a, R.const(alpha, "float32")
                        )
                        lv1: R.Tensor((3, 4), dtype="float32") = R.permute_dims(lv, axes=[1, 0])
                        lv2: R.Tensor((4, 5), dtype="float32") = R.permute_dims(b, axes=[1, 0])
                        gv: R.Tensor((3, 5), dtype="float32") = R.matmul(lv1, lv2)
                        R.output(gv)
                    return gv

            return ExpectedScaledA

        if not useC:

            @I.ir_module
            class ExpectedMatmulOnly:
                @R.function
                def main(
                    a: R.Tensor((4, 3), dtype="float32"),
                    b: R.Tensor((5, 4), dtype="float32"),
                ) -> R.Tensor((3, 5), dtype="float32"):
                    R.func_attr({"num_input": 2})
                    with R.dataflow():
                        lv: R.Tensor((3, 4), dtype="float32") = R.permute_dims(a, axes=[1, 0])
                        lv1: R.Tensor((4, 5), dtype="float32") = R.permute_dims(b, axes=[1, 0])
                        gv: R.Tensor((3, 5), dtype="float32") = R.matmul(lv, lv1)
                        R.output(gv)
                    return gv

            return ExpectedMatmulOnly

        if alpha != 1.0 and beta != 1.0:

            @I.ir_module
            class ExpectedScaledAAndC:
                @R.function
                def main(
                    a: R.Tensor((4, 3), dtype="float32"),
                    b: R.Tensor((5, 4), dtype="float32"),
                    c: R.Tensor((1, 5), dtype="float32"),
                ) -> R.Tensor((3, 5), dtype="float32"):
                    R.func_attr({"num_input": 3})
                    with R.dataflow():
                        lv: R.Tensor((4, 3), dtype="float32") = R.multiply(
                            a, R.const(alpha, "float32")
                        )
                        lv1: R.Tensor((3, 4), dtype="float32") = R.permute_dims(lv, axes=[1, 0])
                        lv2: R.Tensor((4, 5), dtype="float32") = R.permute_dims(b, axes=[1, 0])
                        lv3: R.Tensor((3, 5), dtype="float32") = R.matmul(lv1, lv2)
                        lv4: R.Tensor((1, 5), dtype="float32") = R.multiply(
                            c, R.const(beta, "float32")
                        )
                        gv: R.Tensor((3, 5), dtype="float32") = R.add(lv3, lv4)
                        R.output(gv)
                    return gv

            return ExpectedScaledAAndC

        if alpha != 1.0:

            @I.ir_module
            class ExpectedScaledAWithC:
                @R.function
                def main(
                    a: R.Tensor((4, 3), dtype="float32"),
                    b: R.Tensor((5, 4), dtype="float32"),
                    c: R.Tensor((1, 5), dtype="float32"),
                ) -> R.Tensor((3, 5), dtype="float32"):
                    R.func_attr({"num_input": 3})
                    with R.dataflow():
                        lv: R.Tensor((4, 3), dtype="float32") = R.multiply(
                            a, R.const(alpha, "float32")
                        )
                        lv1: R.Tensor((3, 4), dtype="float32") = R.permute_dims(lv, axes=[1, 0])
                        lv2: R.Tensor((4, 5), dtype="float32") = R.permute_dims(b, axes=[1, 0])
                        lv3: R.Tensor((3, 5), dtype="float32") = R.matmul(lv1, lv2)
                        gv: R.Tensor((3, 5), dtype="float32") = R.add(lv3, c)
                        R.output(gv)
                    return gv

            return ExpectedScaledAWithC

        if beta != 1.0:

            @I.ir_module
            class ExpectedScaledC:
                @R.function
                def main(
                    a: R.Tensor((4, 3), dtype="float32"),
                    b: R.Tensor((5, 4), dtype="float32"),
                    c: R.Tensor((1, 5), dtype="float32"),
                ) -> R.Tensor((3, 5), dtype="float32"):
                    R.func_attr({"num_input": 3})
                    with R.dataflow():
                        lv: R.Tensor((3, 4), dtype="float32") = R.permute_dims(a, axes=[1, 0])
                        lv1: R.Tensor((4, 5), dtype="float32") = R.permute_dims(b, axes=[1, 0])
                        lv2: R.Tensor((3, 5), dtype="float32") = R.matmul(lv, lv1)
                        lv3: R.Tensor((1, 5), dtype="float32") = R.multiply(
                            c, R.const(beta, "float32")
                        )
                        gv: R.Tensor((3, 5), dtype="float32") = R.add(lv2, lv3)
                        R.output(gv)
                    return gv

            return ExpectedScaledC

        @I.ir_module
        class ExpectedMatmulAddC:
            @R.function
            def main(
                a: R.Tensor((4, 3), dtype="float32"),
                b: R.Tensor((5, 4), dtype="float32"),
                c: R.Tensor((1, 5), dtype="float32"),
            ) -> R.Tensor((3, 5), dtype="float32"):
                R.func_attr({"num_input": 3})
                with R.dataflow():
                    lv: R.Tensor((3, 4), dtype="float32") = R.permute_dims(a, axes=[1, 0])
                    lv1: R.Tensor((4, 5), dtype="float32") = R.permute_dims(b, axes=[1, 0])
                    lv2: R.Tensor((3, 5), dtype="float32") = R.matmul(lv, lv1)
                    gv: R.Tensor((3, 5), dtype="float32") = R.add(lv2, c)
                    R.output(gv)
                return gv

        return ExpectedMatmulAddC

    for alpha, beta, useC in [
        (None, None, False),
        (0.25, None, False),
        (1.0, None, False),
        (None, 0.35, False),
        (0.25, 0.35, False),
        (1.0, 0.35, False),
        (None, 1.0, False),
        (0.25, 1.0, False),
        (1.0, 1.0, False),
        (None, None, True),
        (None, 0.35, True),
        (None, 1.0, True),
        (1.0, None, True),
        (1.0, 0.35, True),
        (1.0, 1.0, True),
        (0.25, None, True),
        (0.25, 0.35, True),
        (0.25, 1.0, True),
    ]:
        verify_gemm(alpha, beta, useC, make_expected(alpha, beta, useC))


def test_reshape():
    def verify_reshape(in_shape, shape, out_shape, expected):
        reshape_node = helper.make_node("Reshape", ["data", "shape"], ["reshaped"])

        graph = helper.make_graph(
            [reshape_node],
            "reshape_test",
            inputs=[
                helper.make_tensor_value_info("data", TensorProto.FLOAT, in_shape),
            ],
            initializer=[helper.make_tensor("shape", TensorProto.INT64, [len(shape)], shape)],
            outputs=[helper.make_tensor_value_info("reshaped", TensorProto.FLOAT, out_shape)],
        )
        model = helper.make_model(graph, producer_name="reshape_test")
        tvm_model = from_onnx(model, keep_params_in_input=True)
        tvm_model["main"] = tvm_model["main"].without_attr("params")
        tvm.ir.assert_structural_equal(tvm_model, expected)

    @I.ir_module
    class ExpectedStaticShape:
        @R.function
        def main(
            data: R.Tensor((7, 32, 32, 8), dtype="float32"),
            shape: R.Tensor((2,), dtype="int64"),
        ) -> R.Tensor((224, 256), dtype="float32"):
            R.func_attr({"num_input": 1})
            with R.dataflow():
                gv: R.Tensor((224, 256), dtype="float32") = R.reshape(data, R.shape([224, 256]))
                R.output(gv)
            return gv

    @I.ir_module
    class ExpectedInferDim:
        @R.function
        def main(
            data: R.Tensor((7, 32, 32, 8), dtype="float32"),
            shape: R.Tensor((2,), dtype="int64"),
        ) -> R.Tensor((7, 8192), dtype="float32"):
            R.func_attr({"num_input": 1})
            with R.dataflow():
                gv: R.Tensor((7, 8192), dtype="float32") = R.reshape(data, R.shape([7, 8192]))
                R.output(gv)
            return gv

    @I.ir_module
    class ExpectedCopyInputDim:
        @R.function
        def main(
            data: R.Tensor((7, 32, 32, 8), dtype="float32"),
            shape: R.Tensor((4,), dtype="int64"),
        ) -> R.Tensor((7, 32, 32, 8), dtype="float32"):
            R.func_attr({"num_input": 1})
            with R.dataflow():
                gv: R.Tensor((7, 32, 32, 8), dtype="float32") = R.reshape(
                    data, R.shape([7, 32, 32, 8])
                )
                R.output(gv)
            return gv

    verify_reshape([7, 32, 32, 8], [224, 256], [224, 256], ExpectedStaticShape)
    verify_reshape([7, 32, 32, 8], [-1, 8192], [7, 8192], ExpectedInferDim)
    verify_reshape([7, 32, 32, 8], [0, 32, 32, 8], [7, 32, 32, 8], ExpectedCopyInputDim)


def test_reshape_shape_output():
    def verify_reshape_shape_output(target_shape, output_shape, expected):
        shape_node = helper.make_node("Shape", ["data"], ["shape_out"])
        reshape_node = helper.make_node("Reshape", ["shape_out", "target_shape"], ["reshaped"])

        data_shape = [2, 3, 4]

        graph = helper.make_graph(
            [shape_node, reshape_node],
            "reshape_shape_output",
            inputs=[
                helper.make_tensor_value_info("data", TensorProto.FLOAT, data_shape),
            ],
            initializer=[
                helper.make_tensor(
                    "target_shape", TensorProto.INT64, [len(target_shape)], target_shape
                )
            ],
            outputs=[helper.make_tensor_value_info("reshaped", TensorProto.INT64, output_shape)],
        )
        model = helper.make_model(graph, producer_name="reshape_shape_output")
        tvm_model = from_onnx(model, keep_params_in_input=True)
        assert len(tvm_model["main"].attrs["params"]) == 1
        tvm_model["main"] = tvm_model["main"].without_attr("params")
        tvm.ir.assert_structural_equal(tvm_model, expected)

    @I.ir_module
    class ExpectedFlattenShape:
        @R.function
        def main(
            data: R.Tensor((2, 3, 4), dtype="float32"),
            target_shape: R.Tensor((1,), dtype="int64"),
        ) -> R.Shape([2, 3, 4]):
            R.func_attr({"num_input": 1})
            with R.dataflow():
                gv: R.Shape([2, 3, 4]) = R.shape([2, 3, 4])
                R.output(gv)
            return gv

    @I.ir_module
    class ExpectedRank2Shape:
        @R.function
        def main(
            data: R.Tensor((2, 3, 4), dtype="float32"),
            target_shape: R.Tensor((2,), dtype="int64"),
        ) -> R.Tensor((1, 3), dtype="int64"):
            R.func_attr({"num_input": 1})
            with R.dataflow():
                lv: R.Tensor((3,), dtype="int64") = R.shape_to_tensor(R.shape([2, 3, 4]))
                gv: R.Tensor((1, 3), dtype="int64") = R.reshape(lv, R.shape([1, 3]))
                R.output(gv)
            return gv

    @I.ir_module
    class ExpectedRank2ColumnShape:
        @R.function
        def main(
            data: R.Tensor((2, 3, 4), dtype="float32"),
            target_shape: R.Tensor((2,), dtype="int64"),
        ) -> R.Tensor((3, 1), dtype="int64"):
            R.func_attr({"num_input": 1})
            with R.dataflow():
                lv: R.Tensor((3,), dtype="int64") = R.shape_to_tensor(R.shape([2, 3, 4]))
                gv: R.Tensor((3, 1), dtype="int64") = R.reshape(lv, R.shape([3, 1]))
                R.output(gv)
            return gv

    verify_reshape_shape_output([-1], [3], ExpectedFlattenShape)
    verify_reshape_shape_output([1, 3], [1, 3], ExpectedRank2Shape)
    verify_reshape_shape_output([3, 1], [3, 1], ExpectedRank2ColumnShape)


def test_transpose():
    node = helper.make_node("Transpose", ["x"], ["y"], perm=[1, 2, 0])
    graph = helper.make_graph(
        [node],
        "transpose_test",
        inputs=[helper.make_tensor_value_info("x", TensorProto.FLOAT, [32, 32, 32])],
        outputs=[helper.make_tensor_value_info("y", TensorProto.FLOAT, [32, 32, 32])],
    )
    model = helper.make_model(graph, producer_name="transpose_test")
    check_correctness(model)


def test_transpose_scalar():
    """Test Transpose with scalar inputs - should return scalar unchanged."""
    scalar_node = helper.make_node("Transpose", ["x"], ["y"])
    graph = helper.make_graph(
        [scalar_node],
        "transpose_scalar_test",
        inputs=[helper.make_tensor_value_info("x", TensorProto.FLOAT, [])],
        outputs=[helper.make_tensor_value_info("y", TensorProto.FLOAT, [])],
    )
    model = helper.make_model(graph, producer_name="transpose_scalar_test")
    tvm_model = from_onnx(model, keep_params_in_input=True)

    @I.ir_module
    class ExpectedScalar:
        @R.function
        def main(x: R.Tensor((), dtype="float32")) -> R.Tensor((), dtype="float32"):
            R.func_attr({"num_input": 1})
            with R.dataflow():
                gv: R.Tensor((), dtype="float32") = x
                R.output(gv)
            return gv

    tvm.ir.assert_structural_equal(tvm_model, ExpectedScalar)

    scalar_constant = helper.make_node(
        "Constant",
        [],
        ["scalar"],
        value=helper.make_tensor("value", TensorProto.FLOAT, [], [5.0]),
    )

    transpose_node = helper.make_node("Transpose", ["scalar"], ["y"])
    graph = helper.make_graph(
        [scalar_constant, transpose_node],
        "transpose_scalar_constant_test",
        inputs=[],
        outputs=[helper.make_tensor_value_info("y", TensorProto.FLOAT, [])],
    )
    model = helper.make_model(graph, producer_name="transpose_scalar_constant_test")
    tvm_model = from_onnx(model, keep_params_in_input=True)

    @I.ir_module
    class ExpectedConstant:
        @R.function
        def main() -> R.Tensor((), dtype="float32"):
            R.func_attr({"num_input": 0})
            with R.dataflow():
                gv: R.Tensor((), dtype="float32") = R.const(5.0, "float32")
                R.output(gv)
            return gv

    tvm.ir.assert_structural_equal(tvm_model, ExpectedConstant)


def test_transpose_axes_validation():
    """Test Transpose validation - perm axes count must match tensor dimensions"""

    def assert_transpose_ir(input_shape, axes, output_shape, name, expected):
        transpose_node = helper.make_node("Transpose", ["x"], ["y"], perm=axes)
        graph = helper.make_graph(
            [transpose_node],
            name,
            inputs=[helper.make_tensor_value_info("x", TensorProto.FLOAT, input_shape)],
            outputs=[helper.make_tensor_value_info("y", TensorProto.FLOAT, output_shape)],
        )
        model = helper.make_model(graph, producer_name=name)
        tvm_model = from_onnx(model, keep_params_in_input=True)
        tvm.ir.assert_structural_equal(tvm_model, expected)

    @I.ir_module
    class ExpectedTranspose1D:
        @R.function
        def main(
            x: R.Tensor((10,), dtype="float32"),
        ) -> R.Tensor((10,), dtype="float32"):
            R.func_attr({"num_input": 1})
            with R.dataflow():
                gv: R.Tensor((10,), dtype="float32") = R.permute_dims(x, axes=[0])
                R.output(gv)
            return gv

    @I.ir_module
    class ExpectedTranspose2D:
        @R.function
        def main(
            x: R.Tensor((3, 4), dtype="float32"),
        ) -> R.Tensor((4, 3), dtype="float32"):
            R.func_attr({"num_input": 1})
            with R.dataflow():
                gv: R.Tensor((4, 3), dtype="float32") = R.permute_dims(x, axes=[1, 0])
                R.output(gv)
            return gv

    @I.ir_module
    class ExpectedTranspose3D:
        @R.function
        def main(
            x: R.Tensor((2, 3, 4), dtype="float32"),
        ) -> R.Tensor((4, 2, 3), dtype="float32"):
            R.func_attr({"num_input": 1})
            with R.dataflow():
                gv: R.Tensor((4, 2, 3), dtype="float32") = R.permute_dims(x, axes=[2, 0, 1])
                R.output(gv)
            return gv

    assert_transpose_ir([10], [0], [10], "transpose_1d_valid_test", ExpectedTranspose1D)
    assert_transpose_ir([3, 4], [1, 0], [4, 3], "transpose_2d_valid_test", ExpectedTranspose2D)
    assert_transpose_ir(
        [2, 3, 4], [2, 0, 1], [4, 2, 3], "transpose_3d_valid_test", ExpectedTranspose3D
    )


def assert_static_unsqueeze_ir(
    model: ModelProto,
    *,
    opset: int,
    axes_as_param: bool,
    expected,
):
    tvm_model = from_onnx(model, opset=opset, keep_params_in_input=True)
    if axes_as_param:
        tvm_model["main"] = tvm_model["main"].without_attr("params")
    tvm.ir.assert_structural_equal(tvm_model, expected)


def test_unsqueeze():
    axes = [0, 2, 3]
    unsqueeze_node = helper.make_node("Unsqueeze", ["a", "axes"], ["b"])
    graph = helper.make_graph(
        [unsqueeze_node],
        "unsqueeze",
        inputs=[helper.make_tensor_value_info("a", TensorProto.FLOAT, [32, 32])],
        initializer=[helper.make_tensor("axes", TensorProto.INT64, [3], vals=axes)],
        outputs=[helper.make_tensor_value_info("b", TensorProto.FLOAT, [1, 32, 1, 1, 32])],
    )

    model = helper.make_model(
        graph, producer_name="unsqueeze_test", opset_imports=[helper.make_opsetid("", 13)]
    )

    @I.ir_module
    class ExpectedAxesParam:
        @R.function
        def main(
            a: R.Tensor((32, 32), dtype="float32"),
            axes_param: R.Tensor((3,), dtype="int64"),
        ) -> R.Tensor((1, 32, 1, 1, 32), dtype="float32"):
            R.func_attr({"num_input": 1})
            with R.dataflow():
                lv0: R.Tensor((1, 32, 32), dtype="float32") = R.expand_dims(a, axis=0)
                lv1: R.Tensor((1, 32, 1, 32), dtype="float32") = R.expand_dims(lv0, axis=2)
                gv: R.Tensor((1, 32, 1, 1, 32), dtype="float32") = R.expand_dims(lv1, axis=3)
                R.output(gv)
            return gv

    assert_static_unsqueeze_ir(
        model,
        opset=13,
        axes_as_param=True,
        expected=ExpectedAxesParam,
    )


def test_unsqueeze_scalar_input():
    axes = [0, 1]
    unsqueeze_node = helper.make_node("Unsqueeze", ["a", "axes"], ["b"])

    graph = helper.make_graph(
        [unsqueeze_node],
        "unsqueeze_scalar_input",
        inputs=[helper.make_tensor_value_info("a", TensorProto.FLOAT, [])],
        initializer=[helper.make_tensor("axes", TensorProto.INT64, [2], vals=axes)],
        outputs=[helper.make_tensor_value_info("b", TensorProto.FLOAT, [1, 1])],
    )

    model = helper.make_model(
        graph,
        producer_name="unsqueeze_scalar_input_test",
        opset_imports=[helper.make_opsetid("", 13)],
    )

    @I.ir_module
    class ExpectedScalar:
        @R.function
        def main(
            a: R.Tensor((), dtype="float32"),
            axes_param: R.Tensor((2,), dtype="int64"),
        ) -> R.Tensor((1, 1), dtype="float32"):
            R.func_attr({"num_input": 1})
            with R.dataflow():
                lv0: R.Tensor((1,), dtype="float32") = R.expand_dims(a, axis=0)
                gv: R.Tensor((1, 1), dtype="float32") = R.expand_dims(lv0, axis=1)
                R.output(gv)
            return gv

    assert_static_unsqueeze_ir(
        model,
        opset=13,
        axes_as_param=True,
        expected=ExpectedScalar,
    )


def test_unsqueeze_dynamic_axes_ir():
    unsqueeze_node = helper.make_node("Unsqueeze", ["a", "axes"], ["b"])

    graph = helper.make_graph(
        [unsqueeze_node],
        "unsqueeze_dynamic_axes_ir",
        inputs=[
            helper.make_tensor_value_info("a", TensorProto.FLOAT, [32, 32]),
            helper.make_tensor_value_info("axes", TensorProto.INT64, [2]),
        ],
        outputs=[helper.make_tensor_value_info("b", TensorProto.FLOAT, [1, 32, 32, 1])],
    )

    model = helper.make_model(graph, producer_name="unsqueeze_dynamic_axes_ir_test")
    tvm_model = from_onnx(model, opset=13, keep_params_in_input=True)

    @I.ir_module
    class Expected:
        @R.function
        def main(
            a: R.Tensor((32, 32), dtype="float32"),
            axes: R.Tensor((2,), dtype="int64"),
        ) -> R.Tensor(dtype="float32", ndim=4):
            R.func_attr({"num_input": 2})
            unsqueeze_dim_0 = T.int64()
            unsqueeze_dim_1 = T.int64()
            unsqueeze_dim_2 = T.int64()
            unsqueeze_dim_3 = T.int64()
            with R.dataflow():
                lv: R.Shape([32, 32]) = R.shape_of(a)
                lv1: R.Tensor((2,), dtype="bool") = R.less(axes, R.const(0, "int64"))
                lv2: R.Tensor((2,), dtype="int64") = R.add(axes, R.const(4, "int64"))
                lv3: R.Tensor((4,), dtype="int64") = R.arange(
                    R.prim_value(0), R.prim_value(4), R.prim_value(1), dtype="int64"
                )
                lv4: R.Tensor((2,), dtype="int64") = R.where(lv1, lv2, axes)
                lv5: R.Tensor((4, 1), dtype="int64") = R.expand_dims(lv3, axis=[1])
                lv6: R.Tensor((1, 2), dtype="int64") = R.expand_dims(lv4, axis=[0])
                lv7: R.Tensor((4, 2), dtype="bool") = R.equal(lv5, lv6)
                lv8: R.Tensor((4, 2), dtype="int64") = R.astype(lv7, dtype="int64")
                lv9: R.Tensor((4,), dtype="int64") = R.sum(lv8, axis=[1], keepdims=False)
                lv10: R.Tensor((4,), dtype="int64") = R.subtract(R.const(1, "int64"), lv9)
                lv11: R.Tensor((4,), dtype="int64") = R.cumsum(lv10, axis=0, exclusive=False)
                lv12: R.Tensor((4,), dtype="int64") = R.subtract(lv11, R.const(1, "int64"))
                lv13: R.Tensor((4,), dtype="bool") = R.less(lv12, R.const(0, "int64"))
                lv14: R.Tensor((2,), dtype="int64") = R.shape_to_tensor(lv)
                lv15: R.Tensor((4,), dtype="int64") = R.where(lv13, R.const(0, "int64"), lv12)
                lv16: R.Tensor((4,), dtype="bool") = R.greater(lv9, R.const(0, "int64"))
                lv17: R.Tensor((4,), dtype="int64") = R.take(lv14, lv15, axis=0, mode="fast")
                lv18: R.Tensor((4,), dtype="int64") = R.match_cast(
                    R.where(lv16, R.const(1, "int64"), lv17), R.Tensor((4,), dtype="int64")
                )
                lv19: R.Shape(ndim=4) = R.tensor_to_shape(lv18)
                lv20: R.Shape(
                    [unsqueeze_dim_0, unsqueeze_dim_1, unsqueeze_dim_2, unsqueeze_dim_3]
                ) = R.match_cast(
                    lv19,
                    R.Shape([unsqueeze_dim_0, unsqueeze_dim_1, unsqueeze_dim_2, unsqueeze_dim_3]),
                )
                gv: R.Tensor(
                    (unsqueeze_dim_0, unsqueeze_dim_1, unsqueeze_dim_2, unsqueeze_dim_3),
                    dtype="float32",
                ) = R.reshape(
                    a,
                    R.shape([unsqueeze_dim_0, unsqueeze_dim_1, unsqueeze_dim_2, unsqueeze_dim_3]),
                )
                R.output(gv)
            return gv

    tvm.ir.assert_structural_equal(tvm_model, Expected)


def test_unsqueeze_dynamic_axes_rank_validation():
    unsqueeze_node = helper.make_node("Unsqueeze", ["a", "axes"], ["b"])

    graph = helper.make_graph(
        [unsqueeze_node],
        "unsqueeze_dynamic_axes_rank_validation",
        inputs=[
            helper.make_tensor_value_info("a", TensorProto.FLOAT, [32, 32]),
            helper.make_tensor_value_info("axes", TensorProto.INT64, [1, 2]),
        ],
        outputs=[helper.make_tensor_value_info("b", TensorProto.FLOAT, [1, 32, 32, 1])],
    )

    model = helper.make_model(graph, producer_name="unsqueeze_dynamic_axes_rank_validation_test")
    with pytest.raises(ValueError, match="Expected a 1-D tensor"):
        from_onnx(model, opset=13, keep_params_in_input=True)


def test_unsqueeze_duplicate_axes_validation():
    unsqueeze_node = helper.make_node("Unsqueeze", ["a", "axes"], ["b"])

    graph = helper.make_graph(
        [unsqueeze_node],
        "unsqueeze_duplicate_axes_validation",
        inputs=[helper.make_tensor_value_info("a", TensorProto.FLOAT, [32, 32])],
        initializer=[helper.make_tensor("axes", TensorProto.INT64, [2], vals=[0, 0])],
        outputs=[helper.make_tensor_value_info("b", TensorProto.FLOAT, [1, 1, 32, 32])],
    )

    model = helper.make_model(graph, producer_name="unsqueeze_duplicate_axes_validation_test")
    with pytest.raises(ValueError, match="axes must be unique"):
        from_onnx(model, opset=13)


def test_unsqueeze_v1():
    # https://github.com/onnx/onnx/blob/main/docs/Changelog.md#Unsqueeze-1
    axes = [0, 2, 3]
    unsqueeze_node = helper.make_node("Unsqueeze", ["a"], ["b"], axes=axes)
    graph = helper.make_graph(
        [unsqueeze_node],
        "unsqueeze_v1",
        inputs=[helper.make_tensor_value_info("a", TensorProto.FLOAT, [32, 32])],
        outputs=[helper.make_tensor_value_info("b", TensorProto.FLOAT, [1, 32, 1, 1, 32])],
    )

    model = helper.make_model(
        graph, producer_name="unsqueeze_v1_test", opset_imports=[helper.make_opsetid("", 6)]
    )

    @I.ir_module
    class ExpectedAxesAttr:
        @R.function
        def main(
            a: R.Tensor((32, 32), dtype="float32"),
        ) -> R.Tensor((1, 32, 1, 1, 32), dtype="float32"):
            R.func_attr({"num_input": 1})
            with R.dataflow():
                lv0: R.Tensor((1, 32, 32), dtype="float32") = R.expand_dims(a, axis=0)
                lv1: R.Tensor((1, 32, 1, 32), dtype="float32") = R.expand_dims(lv0, axis=2)
                gv: R.Tensor((1, 32, 1, 1, 32), dtype="float32") = R.expand_dims(lv1, axis=3)
                R.output(gv)
            return gv

    assert_static_unsqueeze_ir(
        model,
        opset=10,
        axes_as_param=False,
        expected=ExpectedAxesAttr,
    )


def test_gelu():
    verify_unary("Gelu", [32, 32], domain="com.microsoft")


def test_gelu_approximate():
    """Test Gelu with approximate attribute from ONNX Opset 20."""
    # Test Gelu with approximate="tanh"
    verify_unary("Gelu", [32, 32], attrs={"approximate": "tanh"}, opset=20)
    # Test Gelu with approximate="none" (default, same as standard Gelu)
    verify_unary("Gelu", [32, 32], attrs={"approximate": "none"}, opset=20)


def test_bias_gelu():
    bias_gelu_node = helper.make_node("BiasGelu", ["a", "b"], ["c"], domain="com.microsoft")
    graph = helper.make_graph(
        [bias_gelu_node],
        "bias_gelu_structural_test",
        inputs=[
            helper.make_tensor_value_info("a", TensorProto.FLOAT, [2, 3]),
            helper.make_tensor_value_info("b", TensorProto.FLOAT, [3]),
        ],
        outputs=[helper.make_tensor_value_info("c", TensorProto.FLOAT, [2, 3])],
    )
    model = helper.make_model(graph, producer_name="bias_gelu_structural_test")
    tvm_model = from_onnx(model, keep_params_in_input=True)

    @I.ir_module
    class Expected:
        @R.function
        def main(
            a: R.Tensor((2, 3), dtype="float32"),
            b: R.Tensor((3,), dtype="float32"),
        ) -> R.Tensor((2, 3), dtype="float32"):
            R.func_attr({"num_input": 2})
            with R.dataflow():
                lv: R.Tensor((2, 3), dtype="float32") = R.add(a, b)
                gv: R.Tensor((2, 3), dtype="float32") = R.nn.gelu(lv)
                R.output(gv)
            return gv

    tvm.ir.assert_structural_equal(tvm_model, Expected)


def test_fast_gelu():
    """Test FastGelu with and without bias"""
    fast_gelu_node = helper.make_node("FastGelu", ["x"], ["y"], domain="com.microsoft")
    graph = helper.make_graph(
        [fast_gelu_node],
        "fast_gelu_structural_test",
        inputs=[helper.make_tensor_value_info("x", TensorProto.FLOAT, [2, 3])],
        outputs=[helper.make_tensor_value_info("y", TensorProto.FLOAT, [2, 3])],
    )
    model = helper.make_model(graph, producer_name="fast_gelu_structural_test")
    tvm_model = from_onnx(model, keep_params_in_input=True)

    @I.ir_module
    class Expected:
        @R.function
        def main(x: R.Tensor((2, 3), dtype="float32")) -> R.Tensor((2, 3), dtype="float32"):
            R.func_attr({"num_input": 1})
            with R.dataflow():
                lv: R.Tensor((2, 3), dtype="float32") = R.multiply(R.const(0.5, "float32"), x)
                lv1: R.Tensor((2, 3), dtype="float32") = R.multiply(
                    R.const(0.79788458347320557, "float32"), x
                )
                lv2: R.Tensor((2, 3), dtype="float32") = R.multiply(x, x)
                lv3: R.Tensor((2, 3), dtype="float32") = R.multiply(lv2, x)
                lv4: R.Tensor((2, 3), dtype="float32") = R.multiply(
                    R.const(0.035677406936883926, "float32"), lv3
                )
                lv5: R.Tensor((2, 3), dtype="float32") = R.add(lv1, lv4)
                lv6: R.Tensor((2, 3), dtype="float32") = R.tanh(lv5)
                lv7: R.Tensor((2, 3), dtype="float32") = R.add(R.const(1.0, "float32"), lv6)
                lv8: R.Tensor((2, 3), dtype="float32") = R.multiply(lv, lv7)
                gv: R.Tensor((2, 3), dtype="float32") = lv8
                R.output(gv)
            return gv

    tvm.ir.assert_structural_equal(tvm_model, Expected)

    fast_gelu_with_bias_node = helper.make_node(
        "FastGelu", ["x", "bias"], ["y"], domain="com.microsoft"
    )
    graph_with_bias = helper.make_graph(
        [fast_gelu_with_bias_node],
        "fast_gelu_with_bias_structural_test",
        inputs=[
            helper.make_tensor_value_info("x", TensorProto.FLOAT, [2, 3]),
            helper.make_tensor_value_info("bias", TensorProto.FLOAT, [3]),
        ],
        outputs=[helper.make_tensor_value_info("y", TensorProto.FLOAT, [2, 3])],
    )
    model_with_bias = helper.make_model(
        graph_with_bias, producer_name="fast_gelu_with_bias_structural_test"
    )
    tvm_model_with_bias = from_onnx(model_with_bias, keep_params_in_input=True)

    @I.ir_module
    class ExpectedWithBias:
        @R.function
        def main(
            x: R.Tensor((2, 3), dtype="float32"),
            bias: R.Tensor((3,), dtype="float32"),
        ) -> R.Tensor((2, 3), dtype="float32"):
            R.func_attr({"num_input": 2})
            with R.dataflow():
                lv: R.Tensor((2, 3), dtype="float32") = R.add(x, bias)
                lv1: R.Tensor((2, 3), dtype="float32") = R.multiply(R.const(0.5, "float32"), lv)
                lv2: R.Tensor((2, 3), dtype="float32") = R.multiply(
                    R.const(0.79788458347320557, "float32"), lv
                )
                lv3: R.Tensor((2, 3), dtype="float32") = R.multiply(lv, lv)
                lv4: R.Tensor((2, 3), dtype="float32") = R.multiply(lv3, lv)
                lv5: R.Tensor((2, 3), dtype="float32") = R.multiply(
                    R.const(0.035677406936883926, "float32"), lv4
                )
                lv6: R.Tensor((2, 3), dtype="float32") = R.add(lv2, lv5)
                lv7: R.Tensor((2, 3), dtype="float32") = R.tanh(lv6)
                lv8: R.Tensor((2, 3), dtype="float32") = R.add(R.const(1.0, "float32"), lv7)
                lv9: R.Tensor((2, 3), dtype="float32") = R.multiply(lv1, lv8)
                gv: R.Tensor((2, 3), dtype="float32") = lv9
                R.output(gv)
            return gv

    tvm.ir.assert_structural_equal(tvm_model_with_bias, ExpectedWithBias)


def test_where():
    where_node = helper.make_node("Where", ["a", "b", "c"], ["d"])

    graph = helper.make_graph(
        [where_node],
        "where_test",
        inputs=[
            helper.make_tensor_value_info("a", TensorProto.BOOL, [32, 32]),
            helper.make_tensor_value_info("b", TensorProto.FLOAT, [32, 32]),
            helper.make_tensor_value_info("c", TensorProto.FLOAT, [32, 32]),
        ],
        outputs=[helper.make_tensor_value_info("d", TensorProto.FLOAT, [32, 32])],
    )

    model = helper.make_model(graph, producer_name="where_test")
    check_correctness(model)


def test_clip():
    def verify_clip(input_names, extra_inputs, expected, tir_func_names=()):
        clip_node = helper.make_node("Clip", input_names, ["output"])
        inputs = [helper.make_tensor_value_info("input", TensorProto.FLOAT, [32, 64])]
        inputs.extend(extra_inputs)
        graph = helper.make_graph(
            [clip_node],
            "clip_test",
            inputs=inputs,
            outputs=[helper.make_tensor_value_info("output", TensorProto.FLOAT, [32, 64])],
        )

        model = helper.make_model(graph, producer_name="clip_test")
        model.opset_import[0].version = 14
        tvm_model = from_onnx(model, keep_params_in_input=True)
        if tir_func_names:
            expected = tvm.IRModule(expected.functions)
            for name in tir_func_names:
                expected.update_func(expected.get_global_var(name), tvm_model[name])
        tvm.ir.assert_structural_equal(tvm_model, expected)

    @I.ir_module
    class ExpectedClipMinMax:
        @T.prim_func(private=True, s_tir=True)
        def maximum(var_input: T.handle, var_min: T.handle, var_output: T.handle):
            T.evaluate(0)

        @T.prim_func(private=True, s_tir=True)
        def minimum(var_input: T.handle, var_max: T.handle, var_output: T.handle):
            T.evaluate(0)

        @R.function
        def main(
            input: R.Tensor((32, 64), dtype="float32"),
            min: R.Tensor((), dtype="float32"),
            max: R.Tensor((), dtype="float32"),
        ) -> R.Tensor((32, 64), dtype="float32"):
            R.func_attr({"num_input": 3})
            cls = ExpectedClipMinMax
            with R.dataflow():
                lv: R.Tensor((), dtype="bool") = R.isnan(min)
                lv1: R.Tensor((), dtype="float32") = R.where(
                    lv, R.const(float("-inf"), "float32"), min
                )
                lv2 = R.call_tir(
                    cls.maximum,
                    (input, lv1),
                    out_ty=R.Tensor((32, 64), dtype="float32"),
                )
                lv3: R.Tensor((), dtype="bool") = R.isnan(max)
                lv4: R.Tensor((), dtype="float32") = R.where(
                    lv3, R.const(float("inf"), "float32"), max
                )
                lv5 = R.call_tir(
                    cls.minimum,
                    (lv2, lv4),
                    out_ty=R.Tensor((32, 64), dtype="float32"),
                )
                gv: R.Tensor((32, 64), dtype="float32") = lv5
                R.output(gv)
            return gv

    @I.ir_module
    class ExpectedClipMin:
        @T.prim_func(private=True, s_tir=True)
        def maximum(var_input: T.handle, var_min: T.handle, var_output: T.handle):
            T.evaluate(0)

        @R.function
        def main(
            input: R.Tensor((32, 64), dtype="float32"),
            min: R.Tensor((), dtype="float32"),
        ) -> R.Tensor((32, 64), dtype="float32"):
            R.func_attr({"num_input": 2})
            cls = ExpectedClipMin
            with R.dataflow():
                lv: R.Tensor((), dtype="bool") = R.isnan(min)
                lv1: R.Tensor((), dtype="float32") = R.where(
                    lv, R.const(float("-inf"), "float32"), min
                )
                lv2 = R.call_tir(
                    cls.maximum,
                    (input, lv1),
                    out_ty=R.Tensor((32, 64), dtype="float32"),
                )
                gv: R.Tensor((32, 64), dtype="float32") = lv2
                R.output(gv)
            return gv

    @I.ir_module
    class ExpectedClipMaxOnlyInput:
        @T.prim_func(private=True, s_tir=True)
        def maximum(var_input: T.handle, var_min: T.handle, var_output: T.handle):
            T.evaluate(0)

        @R.function
        def main(
            input: R.Tensor((32, 64), dtype="float32"),
            max: R.Tensor((), dtype="float32"),
        ) -> R.Tensor((32, 64), dtype="float32"):
            R.func_attr({"num_input": 2})
            cls = ExpectedClipMaxOnlyInput
            with R.dataflow():
                lv: R.Tensor((), dtype="bool") = R.isnan(max)
                lv1: R.Tensor((), dtype="float32") = R.where(
                    lv, R.const(float("-inf"), "float32"), max
                )
                lv2 = R.call_tir(
                    cls.maximum,
                    (input, lv1),
                    out_ty=R.Tensor((32, 64), dtype="float32"),
                )
                gv: R.Tensor((32, 64), dtype="float32") = lv2
                R.output(gv)
            return gv

    @I.ir_module
    class ExpectedClipIdentity:
        @R.function
        def main(
            input: R.Tensor((32, 64), dtype="float32"),
        ) -> R.Tensor((32, 64), dtype="float32"):
            R.func_attr({"num_input": 1})
            with R.dataflow():
                gv: R.Tensor((32, 64), dtype="float32") = input
                R.output(gv)
            return gv

    min_info = helper.make_tensor_value_info("min", TensorProto.FLOAT, ())
    max_info = helper.make_tensor_value_info("max", TensorProto.FLOAT, ())
    verify_clip(
        ["input", "min", "max"], [min_info, max_info], ExpectedClipMinMax, ("maximum", "minimum")
    )
    verify_clip(["input", "min"], [min_info], ExpectedClipMin, ("maximum",))
    verify_clip(["input", "max"], [max_info], ExpectedClipMaxOnlyInput, ("maximum",))
    verify_clip(["input"], [], ExpectedClipIdentity)


@pytest.mark.parametrize("min", [-6.0, 0.0])
@pytest.mark.parametrize("max", [6.0])
def test_clip_v6(max, min):
    # https://github.com/onnx/onnx/blob/main/docs/Changelog.md#Clip-6
    clip_node = helper.make_node("Clip", ["input"], ["output"], max=max, min=min)
    inputs = [helper.make_tensor_value_info("input", TensorProto.FLOAT, [32, 64])]
    graph = helper.make_graph(
        [clip_node],
        "clip_v6_test",
        inputs=inputs,
        outputs=[helper.make_tensor_value_info("output", TensorProto.FLOAT, [32, 64])],
    )
    model = helper.make_model(
        graph, producer_name="clip_v6_test", opset_imports=[helper.make_opsetid("", 6)]
    )
    tvm_model = from_onnx(model, opset=10, keep_params_in_input=True)

    @I.ir_module
    class ExpectedClipV6:
        @T.prim_func(private=True, s_tir=True)
        def maximum(var_input: T.handle, var_output: T.handle):
            T.evaluate(0)

        @T.prim_func(private=True, s_tir=True)
        def minimum(var_input: T.handle, var_output: T.handle):
            T.evaluate(0)

        @R.function
        def main(input: R.Tensor((32, 64), dtype="float32")) -> R.Tensor((32, 64), dtype="float32"):
            R.func_attr({"num_input": 1})
            cls = ExpectedClipV6
            with R.dataflow():
                lv = R.call_tir(
                    cls.maximum,
                    (input,),
                    out_ty=R.Tensor((32, 64), dtype="float32"),
                )
                lv1 = R.call_tir(
                    cls.minimum,
                    (lv,),
                    out_ty=R.Tensor((32, 64), dtype="float32"),
                )
                gv: R.Tensor((32, 64), dtype="float32") = lv1
                R.output(gv)
            return gv

    expected = tvm.IRModule(ExpectedClipV6.functions)
    expected.update_func(expected.get_global_var("maximum"), tvm_model["maximum"])
    expected.update_func(expected.get_global_var("minimum"), tvm_model["minimum"])
    tvm.ir.assert_structural_equal(tvm_model, expected)


@pytest.mark.parametrize(
    "min,max",
    [
        pytest.param(
            np.array(0.0, dtype=np.float32),
            np.array(6.0, dtype=np.float32),
        ),
        pytest.param(
            np.array(0.0, dtype=np.float32),
            np.array(np.nan, dtype=np.float32),
        ),
        pytest.param(
            np.array(np.nan, dtype=np.float32),
            np.array(6.0, dtype=np.float32),
        ),
        pytest.param(
            np.array(np.nan, dtype=np.float32),
            np.array(np.nan, dtype=np.float32),
        ),
    ],
)
@pytest.mark.parametrize(
    "input",
    [
        np.array([0.5, -3.0, 4.5, 11.0, 7.0], dtype=np.float32),
    ],
)
def test_clip_v13(input, min, max):
    # Opset 13: tensor min/max. NaN bound => unbounded on that side (ORT).
    clip_node = helper.make_node("Clip", ["input", "min", "max"], ["output"])
    graph = helper.make_graph(
        [clip_node],
        "clip_v13_nan_max",
        inputs=[
            helper.make_tensor_value_info("input", TensorProto.FLOAT, [5]),
            helper.make_tensor_value_info("min", TensorProto.FLOAT, []),
            helper.make_tensor_value_info("max", TensorProto.FLOAT, []),
        ],
        outputs=[helper.make_tensor_value_info("output", TensorProto.FLOAT, [5])],
    )
    model = helper.make_model(graph, producer_name="clip_v13_nan_max")
    check_correctness(
        model,
        inputs={"input": input, "min": min, "max": max},
        opset=13,
    )


def test_equal():
    equal_node = helper.make_node("Equal", ["a", "b"], ["output"])

    graph = helper.make_graph(
        [equal_node],
        "equal_test",
        inputs=[
            helper.make_tensor_value_info("a", TensorProto.FLOAT, [32, 32]),
            helper.make_tensor_value_info("b", TensorProto.FLOAT, [32, 32]),
        ],
        outputs=[helper.make_tensor_value_info("output", TensorProto.BOOL, [32, 32])],
    )

    model = helper.make_model(graph, producer_name="equal_test")
    check_correctness(
        model, {"a": np.zeros([32, 32], dtype="float32"), "b": np.zeros([32, 32], dtype="float32")}
    )
    check_correctness(
        model, {"a": np.ones([32, 32], dtype="float32"), "b": np.zeros([32, 32], dtype="float32")}
    )
    check_correctness(model)


def test_shape():
    shape_node = helper.make_node("Shape", ["data"], ["output"])

    graph = helper.make_graph(
        [shape_node],
        "shape_test",
        inputs=[
            helper.make_tensor_value_info("data", TensorProto.FLOAT, [3, 4, 5, 6]),
        ],
        outputs=[helper.make_tensor_value_info("output", TensorProto.INT64, [4])],
    )

    model = helper.make_model(graph, producer_name="shape_test")
    tvm_model = from_onnx(model, keep_params_in_input=True)

    @I.ir_module
    class Expected:
        @R.function
        def main(data: R.Tensor((3, 4, 5, 6), dtype="float32")) -> R.Shape([3, 4, 5, 6]):
            R.func_attr({"num_input": 1})
            with R.dataflow():
                gv: R.Shape([3, 4, 5, 6]) = R.shape([3, 4, 5, 6])
                R.output(gv)
            return gv

    tvm.ir.assert_structural_equal(tvm_model, Expected)


def test_trilu():
    def verify_trilu(upper: bool):
        node = helper.make_node("Trilu", ["x"], ["y"], upper=upper)
        graph = helper.make_graph(
            [node],
            "trilu_test",
            inputs=[helper.make_tensor_value_info("x", TensorProto.FLOAT, [3, 5, 5])],
            outputs=[helper.make_tensor_value_info("y", TensorProto.FLOAT, [3, 5, 5])],
        )
        model = helper.make_model(graph, producer_name="trilu_test")
        check_correctness(model)

    verify_trilu(True)
    verify_trilu(False)


@pytest.mark.parametrize("k_value", [-1, 0, 1])
def test_trilu_with_const_k(k_value: int):
    """test_trilu_with_const_k"""

    input_shape = [2, 3, 3]

    graph = helper.make_graph(
        [
            make_constant_node("k", onnx.TensorProto.INT64, [1], [k_value]),
            helper.make_node("Trilu", inputs=["x", "k"], outputs=["y"]),
        ],
        "trilu_graph",
        inputs=[
            helper.make_tensor_value_info("x", onnx.TensorProto.DOUBLE, input_shape),
        ],
        outputs=[helper.make_tensor_value_info("y", onnx.TensorProto.DOUBLE, input_shape)],
    )

    model = helper.make_model(graph, producer_name="trilu_graph")
    check_correctness(model)


def test_selu():
    model = make_unary_model("Selu", [2, 3])
    tvm_model = from_onnx(model, keep_params_in_input=True)

    @I.ir_module
    class Expected:
        @R.function
        def main(x: R.Tensor((2, 3), dtype="float32")) -> R.Tensor((2, 3), dtype="float32"):
            R.func_attr({"num_input": 1})
            with R.dataflow():
                lv: R.Tensor((2, 3), dtype="float32") = R.exp(x)
                lv1: R.Tensor((2, 3), dtype="float32") = R.subtract(R.const(1.0, "float32"), lv)
                lv2: R.Tensor((2, 3), dtype="float32") = R.nn.relu(lv1)
                lv3: R.Tensor((2, 3), dtype="float32") = R.multiply(
                    R.const(-1.6732631921768188, "float32"), lv2
                )
                lv4: R.Tensor((2, 3), dtype="float32") = R.nn.relu(x)
                lv5: R.Tensor((2, 3), dtype="float32") = R.add(lv3, lv4)
                gv: R.Tensor((2, 3), dtype="float32") = R.multiply(
                    R.const(1.0507010221481323, "float32"), lv5
                )
                R.output(gv)
            return gv

    tvm.ir.assert_structural_equal(tvm_model, Expected)

    model = make_unary_model("Selu", [2, 3], attrs={"alpha": 0.25, "gamma": 0.3})
    tvm_model = from_onnx(model, keep_params_in_input=True)

    @I.ir_module
    class ExpectedCustom:
        @R.function
        def main(x: R.Tensor((2, 3), dtype="float32")) -> R.Tensor((2, 3), dtype="float32"):
            R.func_attr({"num_input": 1})
            with R.dataflow():
                lv: R.Tensor((2, 3), dtype="float32") = R.exp(x)
                lv1: R.Tensor((2, 3), dtype="float32") = R.subtract(R.const(1.0, "float32"), lv)
                lv2: R.Tensor((2, 3), dtype="float32") = R.nn.relu(lv1)
                lv3: R.Tensor((2, 3), dtype="float32") = R.multiply(R.const(-0.25, "float32"), lv2)
                lv4: R.Tensor((2, 3), dtype="float32") = R.nn.relu(x)
                lv5: R.Tensor((2, 3), dtype="float32") = R.add(lv3, lv4)
                gv: R.Tensor((2, 3), dtype="float32") = R.multiply(
                    R.const(0.30000001192092896, "float32"), lv5
                )
                R.output(gv)
            return gv

    tvm.ir.assert_structural_equal(tvm_model, ExpectedCustom)


def test_mish():
    model = make_unary_model("Mish", [2, 3])
    tvm_model = from_onnx(model, opset=18, keep_params_in_input=True)

    @I.ir_module
    class Expected:
        @R.function
        def main(x: R.Tensor((2, 3), dtype="float32")) -> R.Tensor((2, 3), dtype="float32"):
            R.func_attr({"num_input": 1})
            with R.dataflow():
                lv: R.Tensor((2, 3), dtype="float32") = R.exp(x)
                lv1: R.Tensor((2, 3), dtype="float32") = R.add(R.const(1.0, "float32"), lv)
                lv2: R.Tensor((2, 3), dtype="float32") = R.log(lv1)
                lv3: R.Tensor((2, 3), dtype="float32") = R.tanh(lv2)
                gv: R.Tensor((2, 3), dtype="float32") = R.multiply(x, lv3)
                R.output(gv)
            return gv

    tvm.ir.assert_structural_equal(tvm_model, Expected)


def test_prelu():
    def _assert_prelu_ir(slope_shape, expected):
        prelu_node = helper.make_node("PRelu", ["a", "b"], ["c"])
        graph = helper.make_graph(
            [prelu_node],
            "prelu_structural_test",
            inputs=[
                helper.make_tensor_value_info("a", TensorProto.FLOAT, [3, 32, 32]),
                helper.make_tensor_value_info("b", TensorProto.FLOAT, slope_shape),
            ],
            outputs=[helper.make_tensor_value_info("c", TensorProto.FLOAT, [3, 32, 32])],
        )
        model = helper.make_model(graph, producer_name="prelu_structural_test")
        tvm_model = from_onnx(model, keep_params_in_input=True)
        tvm.ir.assert_structural_equal(tvm_model, expected)

    @I.ir_module
    class ExpectedScalarSlope:
        @R.function
        def main(
            a: R.Tensor((3, 32, 32), dtype="float32"),
            b: R.Tensor((1,), dtype="float32"),
        ) -> R.Tensor((3, 32, 32), dtype="float32"):
            R.func_attr({"num_input": 2})
            with R.dataflow():
                lv: R.Tensor((1,), dtype="float32") = R.reshape(b, R.shape([1]))
                gv: R.Tensor((3, 32, 32), dtype="float32") = R.nn.prelu(a, lv, axis=2)
                R.output(gv)
            return gv

    @I.ir_module
    class ExpectedTwoDimScalarSlope:
        @R.function
        def main(
            a: R.Tensor((3, 32, 32), dtype="float32"),
            b: R.Tensor((1, 1), dtype="float32"),
        ) -> R.Tensor((3, 32, 32), dtype="float32"):
            R.func_attr({"num_input": 2})
            with R.dataflow():
                lv: R.Tensor((1,), dtype="float32") = R.reshape(b, R.shape([1]))
                gv: R.Tensor((3, 32, 32), dtype="float32") = R.nn.prelu(a, lv, axis=2)
                R.output(gv)
            return gv

    @I.ir_module
    class ExpectedChannelSlope:
        @R.function
        def main(
            a: R.Tensor((3, 32, 32), dtype="float32"),
            b: R.Tensor((32,), dtype="float32"),
        ) -> R.Tensor((3, 32, 32), dtype="float32"):
            R.func_attr({"num_input": 2})
            with R.dataflow():
                lv: R.Tensor((32,), dtype="float32") = R.reshape(b, R.shape([32]))
                gv: R.Tensor((3, 32, 32), dtype="float32") = R.nn.prelu(a, lv, axis=2)
                R.output(gv)
            return gv

    @I.ir_module
    class ExpectedBatchSlope:
        @R.function
        def main(
            a: R.Tensor((3, 32, 32), dtype="float32"),
            b: R.Tensor((3, 1, 1), dtype="float32"),
        ) -> R.Tensor((3, 32, 32), dtype="float32"):
            R.func_attr({"num_input": 2})
            with R.dataflow():
                lv: R.Tensor((3,), dtype="float32") = R.reshape(b, R.shape([3]))
                gv: R.Tensor((3, 32, 32), dtype="float32") = R.nn.prelu(a, lv, axis=0)
                R.output(gv)
            return gv

    _assert_prelu_ir([1], ExpectedScalarSlope)
    _assert_prelu_ir([1, 1], ExpectedTwoDimScalarSlope)
    _assert_prelu_ir([32], ExpectedChannelSlope)
    _assert_prelu_ir([3, 1, 1], ExpectedBatchSlope)


def test_thresholded_relu():
    model = make_unary_model("ThresholdedRelu", [2, 3])
    tvm_model = from_onnx(model, keep_params_in_input=True)

    @I.ir_module
    class Expected:
        @R.function
        def main(x: R.Tensor((2, 3), dtype="float32")) -> R.Tensor((2, 3), dtype="float32"):
            R.func_attr({"num_input": 1})
            with R.dataflow():
                lv: R.Tensor((2, 3), dtype="bool") = R.greater(x, R.const(1.0, "float32"))
                lv1: R.Tensor((2, 3), dtype="float32") = R.astype(lv, dtype="float32")
                gv: R.Tensor((2, 3), dtype="float32") = R.multiply(lv1, x)
                R.output(gv)
            return gv

    tvm.ir.assert_structural_equal(tvm_model, Expected)

    model = make_unary_model("ThresholdedRelu", [2, 3], attrs={"alpha": -0.01})
    tvm_model = from_onnx(model, keep_params_in_input=True)

    @I.ir_module
    class ExpectedCustom:
        @R.function
        def main(x: R.Tensor((2, 3), dtype="float32")) -> R.Tensor((2, 3), dtype="float32"):
            R.func_attr({"num_input": 1})
            with R.dataflow():
                lv: R.Tensor((2, 3), dtype="bool") = R.greater(
                    x, R.const(-0.0099999997764825821, "float32")
                )
                lv1: R.Tensor((2, 3), dtype="float32") = R.astype(lv, dtype="float32")
                gv: R.Tensor((2, 3), dtype="float32") = R.multiply(lv1, x)
                R.output(gv)
            return gv

    tvm.ir.assert_structural_equal(tvm_model, ExpectedCustom)


def test_leakyrelu():
    verify_unary("LeakyRelu", [32, 32])
    verify_unary("LeakyRelu", [32, 32], attrs={"alpha": 0.2})


def test_hardsigmoid():
    model = make_unary_model("HardSigmoid", [2, 3])
    tvm_model = from_onnx(model, keep_params_in_input=True)

    @I.ir_module
    class Expected:
        @R.function
        def main(x: R.Tensor((2, 3), dtype="float32")) -> R.Tensor((2, 3), dtype="float32"):
            R.func_attr({"num_input": 1})
            with R.dataflow():
                lv: R.Tensor((2, 3), dtype="float32") = R.multiply(
                    R.const(0.20000000298023224, "float32"), x
                )
                lv1: R.Tensor((2, 3), dtype="float32") = R.add(lv, R.const(0.5, "float32"))
                gv: R.Tensor((2, 3), dtype="float32") = R.clip(
                    lv1, R.prim_value(0), R.prim_value(1)
                )
                R.output(gv)
            return gv

    tvm.ir.assert_structural_equal(tvm_model, Expected)

    model = make_unary_model("HardSigmoid", [2, 3], attrs={"alpha": 0.3, "beta": 0.4})
    tvm_model = from_onnx(model, keep_params_in_input=True)

    @I.ir_module
    class ExpectedCustom:
        @R.function
        def main(x: R.Tensor((2, 3), dtype="float32")) -> R.Tensor((2, 3), dtype="float32"):
            R.func_attr({"num_input": 1})
            with R.dataflow():
                lv: R.Tensor((2, 3), dtype="float32") = R.multiply(
                    R.const(0.30000001192092896, "float32"), x
                )
                lv1: R.Tensor((2, 3), dtype="float32") = R.add(
                    lv, R.const(0.40000000596046448, "float32")
                )
                gv: R.Tensor((2, 3), dtype="float32") = R.clip(
                    lv1, R.prim_value(0), R.prim_value(1)
                )
                R.output(gv)
            return gv

    tvm.ir.assert_structural_equal(tvm_model, ExpectedCustom)

    model = make_unary_model("HardSigmoid", [1, 3, 20, 20], attrs={"alpha": 0.5, "beta": 0.6})
    tvm_model = from_onnx(model, keep_params_in_input=True)

    @I.ir_module
    class ExpectedCustom4D:
        @R.function
        def main(
            x: R.Tensor((1, 3, 20, 20), dtype="float32"),
        ) -> R.Tensor((1, 3, 20, 20), dtype="float32"):
            R.func_attr({"num_input": 1})
            with R.dataflow():
                lv: R.Tensor((1, 3, 20, 20), dtype="float32") = R.multiply(
                    R.const(0.5, "float32"), x
                )
                lv1: R.Tensor((1, 3, 20, 20), dtype="float32") = R.add(
                    lv, R.const(0.60000002384185791, "float32")
                )
                gv: R.Tensor((1, 3, 20, 20), dtype="float32") = R.clip(
                    lv1, R.prim_value(0), R.prim_value(1)
                )
                R.output(gv)
            return gv

    tvm.ir.assert_structural_equal(tvm_model, ExpectedCustom4D)


def test_shrink():
    model = make_unary_model("Shrink", [2, 3])
    tvm_model = from_onnx(model, keep_params_in_input=True)

    @I.ir_module
    class Expected:
        @R.function
        def main(x: R.Tensor((2, 3), dtype="float32")) -> R.Tensor((2, 3), dtype="float32"):
            R.func_attr({"num_input": 1})
            with R.dataflow():
                lv: R.Tensor((2, 3), dtype="bool") = R.greater(x, R.const(0.5, "float32"))
                lv1: R.Tensor((2, 3), dtype="float32") = R.subtract(x, R.const(0.0, "float32"))
                lv2: R.Tensor((2, 3), dtype="float32") = R.zeros_like(x)
                lv3: R.Tensor((2, 3), dtype="float32") = R.where(lv, lv1, lv2)
                lv4: R.Tensor((), dtype="float32") = R.negative(R.const(0.5, "float32"))
                lv5: R.Tensor((2, 3), dtype="bool") = R.less(x, lv4)
                lv6: R.Tensor((2, 3), dtype="float32") = R.add(x, R.const(0.0, "float32"))
                lv7: R.Tensor((2, 3), dtype="float32") = R.where(lv5, lv6, lv2)
                gv: R.Tensor((2, 3), dtype="float32") = R.add(lv3, lv7)
                R.output(gv)
            return gv

    tvm.ir.assert_structural_equal(tvm_model, Expected)

    model = make_unary_model("Shrink", [2, 3], attrs={"lambd": 0.2, "bias": 0.1})
    tvm_model = from_onnx(model, keep_params_in_input=True)

    @I.ir_module
    class ExpectedCustom:
        @R.function
        def main(x: R.Tensor((2, 3), dtype="float32")) -> R.Tensor((2, 3), dtype="float32"):
            R.func_attr({"num_input": 1})
            with R.dataflow():
                lv: R.Tensor((2, 3), dtype="bool") = R.greater(
                    x, R.const(0.20000000298023224, "float32")
                )
                lv1: R.Tensor((2, 3), dtype="float32") = R.subtract(
                    x, R.const(0.10000000149011612, "float32")
                )
                lv2: R.Tensor((2, 3), dtype="float32") = R.zeros_like(x)
                lv3: R.Tensor((2, 3), dtype="float32") = R.where(lv, lv1, lv2)
                lv4: R.Tensor((), dtype="float32") = R.negative(
                    R.const(0.20000000298023224, "float32")
                )
                lv5: R.Tensor((2, 3), dtype="bool") = R.less(x, lv4)
                lv6: R.Tensor((2, 3), dtype="float32") = R.add(
                    x, R.const(0.10000000149011612, "float32")
                )
                lv7: R.Tensor((2, 3), dtype="float32") = R.where(lv5, lv6, lv2)
                gv: R.Tensor((2, 3), dtype="float32") = R.add(lv3, lv7)
                R.output(gv)
            return gv

    tvm.ir.assert_structural_equal(tvm_model, ExpectedCustom)


@pytest.mark.parametrize("stride", [1, 2])
@pytest.mark.parametrize("dilation", [1, 2])
@pytest.mark.parametrize("bias", [True, False])
@pytest.mark.parametrize("pad", [0, 2])
@pytest.mark.parametrize("auto_pad", ["SAME_UPPER", "SAME_LOWER", "VALID"])
def test_conv(stride: int, dilation: int, pad: int, bias: bool, auto_pad: str):
    def _verify_conv(input_shape, weight_shape):
        nd = len(weight_shape) - 2
        if auto_pad == "VALID":
            output_shape = [input_shape[0], weight_shape[0]] + [
                (input_shape[i] - dilation * (weight_shape[i] - 1) - 1) // stride + 1
                for i in range(2, len(input_shape))
            ]
            bias_shape = [output_shape[1]]
            conv_node = helper.make_node(
                "Conv",
                inputs=["x", "w"] + (["b"] if bias else []),
                outputs=["y"],
                strides=[stride] * nd,
                dilations=[dilation] * nd,
                auto_pad=auto_pad,
                group=input_shape[1] // weight_shape[1],
            )
        elif auto_pad in ("SAME_UPPER", "SAME_LOWER"):
            if dilation == 2:
                # auto_pad = "SAME" and dilation = 2 is not supported in ONNX
                return
            output_shape = [input_shape[0], weight_shape[0]] + [
                (input_shape[i] + stride - 1) // stride for i in range(2, len(input_shape))
            ]
            bias_shape = [output_shape[1]]
            conv_node = helper.make_node(
                "Conv",
                inputs=["x", "w"] + (["b"] if bias else []),
                outputs=["y"],
                strides=[stride] * nd,
                dilations=[dilation] * nd,
                auto_pad=auto_pad,
                group=input_shape[1] // weight_shape[1],
            )
        else:
            output_shape = [input_shape[0], weight_shape[0]] + [
                (input_shape[i] + 2 * pad - dilation * (weight_shape[i] - 1) - 1) // stride + 1
                for i in range(2, len(input_shape))
            ]
            bias_shape = [output_shape[1]]
            conv_node = helper.make_node(
                "Conv",
                inputs=["x", "w"] + (["b"] if bias else []),
                outputs=["y"],
                strides=[stride] * nd,
                dilations=[dilation] * nd,
                pads=[pad] * nd * 2,
                group=input_shape[1] // weight_shape[1],
            )
        graph = helper.make_graph(
            [conv_node],
            "conv_test",
            inputs=[
                helper.make_tensor_value_info("x", TensorProto.FLOAT, input_shape),
                helper.make_tensor_value_info("w", TensorProto.FLOAT, weight_shape),
            ]
            + ([helper.make_tensor_value_info("b", TensorProto.FLOAT, bias_shape)] if bias else []),
            outputs=[helper.make_tensor_value_info("y", TensorProto.FLOAT, output_shape)],
        )

        model = helper.make_model(graph, producer_name="conv_test")
        check_correctness(model, atol=1e-4)

    # Conv1D
    _verify_conv([3, 4, 32], [4, 4, 3])
    _verify_conv([3, 4, 32], [2, 4, 3])  # group=2
    # Conv2D
    _verify_conv([3, 4, 32, 32], [4, 4, 3, 3])
    _verify_conv([3, 4, 32, 32], [2, 4, 3, 3])  # group=2
    # Conv3D
    _verify_conv([3, 4, 32, 32, 32], [4, 4, 3, 3, 3])
    _verify_conv([3, 4, 32, 32, 32], [2, 4, 3, 3, 3])  # group=2


@pytest.mark.parametrize("stride", [2])
@pytest.mark.parametrize("dilation", [1])
@pytest.mark.parametrize("bias", [True, False])
@pytest.mark.parametrize("pad", [0, 2])
@pytest.mark.parametrize("output_pad", [0, 1])
def test_conv_transpose(stride: int, dilation: int, pad: int, bias: bool, output_pad: int):
    def _verify_conv_transpose(input_shape, weight_shape):
        nd = len(weight_shape) - 2
        output_shape = [input_shape[0], weight_shape[0]] + [
            (input_shape[i] - 1) * stride
            - 2 * pad
            + dilation * (weight_shape[i] - 1)
            + output_pad
            + 1
            for i in range(2, len(input_shape))
        ]
        bias_shape = [output_shape[1]]
        conv_node = helper.make_node(
            "ConvTranspose",
            inputs=["x", "w"] + (["b"] if bias else []),
            outputs=["y"],
            strides=[stride] * nd,
            dilations=[dilation] * nd,
            pads=[pad] * nd * 2,
            output_padding=[output_pad] * nd,
            group=input_shape[1] // weight_shape[1],
        )
        graph = helper.make_graph(
            [conv_node],
            "conv_transpose_test",
            inputs=[
                helper.make_tensor_value_info("x", TensorProto.FLOAT, input_shape),
                helper.make_tensor_value_info("w", TensorProto.FLOAT, weight_shape),
            ]
            + ([helper.make_tensor_value_info("b", TensorProto.FLOAT, bias_shape)] if bias else []),
            outputs=[helper.make_tensor_value_info("y", TensorProto.FLOAT, output_shape)],
        )

        model = helper.make_model(graph, producer_name="conv_transpose_test")
        check_correctness(model, atol=1e-4)

    # ConvTranspose1D
    _verify_conv_transpose([3, 4, 32], [4, 4, 3])
    _verify_conv_transpose([3, 4, 32], [4, 2, 3])  # group=2
    # ConvTranspose2D
    _verify_conv_transpose([3, 4, 32, 32], [4, 4, 3, 3])
    _verify_conv_transpose([3, 4, 32, 32], [4, 2, 3, 3])  # group=2
    # ConvTranspose3D
    _verify_conv_transpose([3, 4, 12, 12, 12], [4, 4, 3, 3, 3])
    _verify_conv_transpose([3, 4, 12, 12, 12], [4, 2, 3, 3, 3])  # group=2


@pytest.mark.parametrize("auto_pad", ["SAME_UPPER", "SAME_LOWER", "VALID"])
@pytest.mark.parametrize("stride", [1, 2])
def test_conv_transpose_auto_pad(auto_pad: str, stride: int):
    def _verify(input_shape, weight_shape):
        nd = len(weight_shape) - 2
        conv_node = helper.make_node(
            "ConvTranspose",
            inputs=["x", "w"],
            outputs=["y"],
            kernel_shape=weight_shape[2:],
            strides=[stride] * nd,
            auto_pad=auto_pad,
        )
        graph = helper.make_graph(
            [conv_node],
            "conv_transpose_auto_pad_test",
            inputs=[
                helper.make_tensor_value_info("x", TensorProto.FLOAT, input_shape),
                helper.make_tensor_value_info("w", TensorProto.FLOAT, weight_shape),
            ],
            outputs=[helper.make_tensor_value_info("y", TensorProto.FLOAT, None)],
        )
        model = helper.make_model(graph, producer_name="conv_transpose_auto_pad_test")
        check_correctness(model, atol=1e-4)

    # ConvTranspose1D / 2D / 3D
    _verify([1, 1, 8], [1, 1, 3])
    _verify([1, 1, 8, 8], [1, 1, 3, 3])
    _verify([1, 1, 4, 4, 4], [1, 1, 3, 3, 3])


def test_pow():
    verify_binary("Pow", [32, 32], [32, 32], [32, 32])


@pytest.mark.parametrize("reverse", [True, False])
@pytest.mark.parametrize("exclusive", [True, False])
def test_cumsum(reverse, exclusive):
    cumsum_node = helper.make_node(
        "CumSum", ["x", "axis"], ["y"], reverse=reverse, exclusive=exclusive
    )
    shape = [32, 32]
    graph = helper.make_graph(
        [cumsum_node],
        "cumsum_test",
        inputs=[
            helper.make_tensor_value_info("x", TensorProto.FLOAT, shape),
        ],
        initializer=[helper.make_tensor("axis", TensorProto.INT64, (), [1])],
        outputs=[helper.make_tensor_value_info("y", TensorProto.FLOAT, shape)],
    )

    model = helper.make_model(
        graph, producer_name="cumsum_test", opset_imports=[helper.make_opsetid("", 14)]
    )
    check_correctness(model)


def test_cumsum_int32_1d_axis_initializer():
    input_shape = [2, 3]

    graph = helper.make_graph(
        [
            helper.make_node("CumSum", inputs=["X", "axis"], outputs=["Y"]),
        ],
        "cumsum_graph",
        inputs=[
            helper.make_tensor_value_info("X", onnx.TensorProto.DOUBLE, input_shape),
        ],
        initializer=[helper.make_tensor("axis", onnx.TensorProto.INT32, [1], [0])],
        outputs=[helper.make_tensor_value_info("Y", onnx.TensorProto.DOUBLE, input_shape)],
    )

    model = helper.make_model(graph, producer_name="cumsum_graph")
    check_correctness(model)


def test_cumsum_dynamic_axis_not_supported():
    input_shape = [2, 3]

    graph = helper.make_graph(
        [
            helper.make_node("CumSum", inputs=["X", "axis"], outputs=["Y"]),
        ],
        "cumsum_dynamic_axis_graph",
        inputs=[
            helper.make_tensor_value_info("X", onnx.TensorProto.DOUBLE, input_shape),
            helper.make_tensor_value_info("axis", onnx.TensorProto.INT32, [1], "axis"),
        ],
        outputs=[helper.make_tensor_value_info("Y", onnx.TensorProto.DOUBLE, input_shape)],
    )

    model = helper.make_model(graph, producer_name="cumsum_dynamic_axis_graph")
    with pytest.raises(ValueError, match="non-constant axis input is not supported"):
        from_onnx(model, opset=14, keep_params_in_input=True)


def test_cumsum_axis_shape_validation():
    input_shape = [2, 3]

    graph = helper.make_graph(
        [
            helper.make_node("CumSum", inputs=["X", "axis"], outputs=["Y"]),
        ],
        "cumsum_invalid_axis_shape_graph",
        inputs=[
            helper.make_tensor_value_info("X", onnx.TensorProto.DOUBLE, input_shape),
        ],
        initializer=[helper.make_tensor("axis", onnx.TensorProto.INT64, [2], [0, 1])],
        outputs=[helper.make_tensor_value_info("Y", onnx.TensorProto.DOUBLE, input_shape)],
    )

    model = helper.make_model(graph, producer_name="cumsum_invalid_axis_shape_graph")
    with pytest.raises(
        ValueError,
        match=r"axis input must be a scalar \(0-D\) or a single-element 1-D tensor",
    ):
        from_onnx(model, opset=14, keep_params_in_input=True)


def test_squeeze():
    def verify_squeeze(axis, expected):
        if axis:
            squeeze_node = helper.make_node("Squeeze", ["x", "axes"], ["y"])
        else:
            squeeze_node = helper.make_node("Squeeze", ["x"], ["y"])
        shape = [1, 32, 1, 32]

        initializer = (
            [helper.make_tensor("axes", TensorProto.INT64, [len(axis)], axis)] if axis else None
        )

        graph = helper.make_graph(
            [squeeze_node],
            "squeeze_test",
            inputs=[
                helper.make_tensor_value_info("x", TensorProto.FLOAT, shape),
            ],
            initializer=initializer,
            outputs=[helper.make_tensor_value_info("y", TensorProto.FLOAT, [32, 32])],
        )

        model = helper.make_model(
            graph, producer_name="squeeze_test", opset_imports=[helper.make_opsetid("", 13)]
        )
        tvm_model = from_onnx(model, opset=13, keep_params_in_input=True)
        if axis:
            tvm_model["main"] = tvm_model["main"].without_attr("params")
        tvm.ir.assert_structural_equal(tvm_model, expected)

    @I.ir_module
    class ExpectedSqueezeAxes:
        @R.function
        def main(
            x: R.Tensor((1, 32, 1, 32), dtype="float32"),
            axes: R.Tensor((2,), dtype="int64"),
        ) -> R.Tensor((32, 32), dtype="float32"):
            R.func_attr({"num_input": 1})
            with R.dataflow():
                gv: R.Tensor((32, 32), dtype="float32") = R.squeeze(x, axis=[0, 2])
                R.output(gv)
            return gv

    @I.ir_module
    class ExpectedSqueezeAll:
        @R.function
        def main(x: R.Tensor((1, 32, 1, 32), dtype="float32")) -> R.Tensor(
            (32, 32), dtype="float32"
        ):
            R.func_attr({"num_input": 1})
            with R.dataflow():
                gv: R.Tensor((32, 32), dtype="float32") = R.squeeze(x, axis=None)
                R.output(gv)
            return gv

    verify_squeeze([0, 2], ExpectedSqueezeAxes)
    verify_squeeze(None, ExpectedSqueezeAll)


def test_squeeze_constant():
    def verify_squeeze_constant(axis, expected):
        shape = [1, 2, 1, 3]
        data = np.arange(6, dtype="float32").reshape(shape)
        constant = make_constant_node("x", onnx.TensorProto.FLOAT, shape, data.flatten().tolist())
        if axis:
            squeeze_node = helper.make_node("Squeeze", ["x", "axes"], ["y"])
        else:
            squeeze_node = helper.make_node("Squeeze", ["x"], ["y"])

        initializer = (
            [helper.make_tensor("axes", TensorProto.INT64, [len(axis)], axis)] if axis else None
        )

        graph = helper.make_graph(
            [constant, squeeze_node],
            "squeeze_test",
            inputs=[],
            initializer=initializer,
            outputs=[helper.make_tensor_value_info("y", TensorProto.FLOAT, [2, 3])],
        )

        model = helper.make_model(
            graph, producer_name="squeeze_test", opset_imports=[helper.make_opsetid("", 13)]
        )
        tvm_model = from_onnx(model, opset=13, keep_params_in_input=True)
        if axis:
            tvm_model["main"] = tvm_model["main"].without_attr("params")
        tvm.ir.assert_structural_equal(tvm_model, expected)

    @I.ir_module
    class ExpectedSqueezeConstantAxes:
        @R.function
        def main(axes: R.Tensor((2,), dtype="int64")) -> R.Tensor((2, 3), dtype="float32"):
            R.func_attr({"num_input": 0})
            with R.dataflow():
                gv: R.Tensor((2, 3), dtype="float32") = R.const(
                    [[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]], "float32"
                )
                R.output(gv)
            return gv

    @I.ir_module
    class ExpectedSqueezeConstantAll:
        @R.function
        def main() -> R.Tensor((2, 3), dtype="float32"):
            R.func_attr({"num_input": 0})
            with R.dataflow():
                gv: R.Tensor((2, 3), dtype="float32") = R.const(
                    [[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]], "float32"
                )
                R.output(gv)
            return gv

    verify_squeeze_constant([0, 2], ExpectedSqueezeConstantAxes)
    verify_squeeze_constant(None, ExpectedSqueezeConstantAll)


@pytest.mark.parametrize("axis", [[0]])
@pytest.mark.parametrize("A", [8, 16, 32])
@pytest.mark.parametrize("B", [8, 16, 32])
def test_dynamic_squeeze(axis, A, B):
    squeeze_node = helper.make_node("Squeeze", ["x", "axes"], ["y"])
    shape = [1, "A", "B"]

    initializer = (
        [helper.make_tensor("axes", TensorProto.INT64, [len(axis)], axis)] if axis else None
    )

    graph = helper.make_graph(
        [squeeze_node],
        "squeeze_test",
        inputs=[
            helper.make_tensor_value_info("x", TensorProto.FLOAT, shape),
        ],
        initializer=initializer,
        outputs=[helper.make_tensor_value_info("y", TensorProto.FLOAT, ["A", "B"])],
    )

    model = helper.make_model(graph, producer_name="squeeze_test")
    tvm_model = from_onnx(model, opset=13, keep_params_in_input=True)
    tvm_model["main"] = tvm_model["main"].without_attr("params")

    @I.ir_module
    class Expected:
        @R.function
        def main(
            x: R.Tensor((1, "A", "B"), dtype="float32"),
            axes: R.Tensor((1,), dtype="int64"),
        ) -> R.Tensor(("A", "B"), dtype="float32"):
            A = T.int64()
            B = T.int64()
            R.func_attr({"num_input": 1})
            with R.dataflow():
                gv: R.Tensor((A, B), dtype="float32") = R.squeeze(x, axis=[0])
                R.output(gv)
            return gv

    tvm.ir.assert_structural_equal(tvm_model, Expected)


def test_squeeze_dynamic_axes_ir():
    squeeze_node = helper.make_node("Squeeze", ["x", "axes"], ["y"])
    shape = [1, 32, 1, 32]

    graph = helper.make_graph(
        [squeeze_node],
        "squeeze_dynamic_axes_ir",
        inputs=[
            helper.make_tensor_value_info("x", TensorProto.FLOAT, shape),
            helper.make_tensor_value_info("axes", TensorProto.INT64, [2]),
        ],
        outputs=[helper.make_tensor_value_info("y", TensorProto.FLOAT, [32, 32])],
    )

    model = helper.make_model(graph, producer_name="squeeze_dynamic_axes_ir_test")
    tvm_model = from_onnx(model, opset=13, keep_params_in_input=True)

    @I.ir_module
    class Expected:
        @R.function
        def main(
            x: R.Tensor((1, 32, 1, 32), dtype="float32"),
            axes: R.Tensor((2,), dtype="int64"),
        ) -> R.Tensor(dtype="float32", ndim=2):
            R.func_attr({"num_input": 2})
            squeeze_num_keep_dims = T.int64()
            squeeze_dim_0 = T.int64()
            squeeze_dim_1 = T.int64()
            with R.dataflow():
                lv: R.Shape([1, 32, 1, 32]) = R.shape_of(x)
                lv1: R.Tensor((2,), dtype="bool") = R.less(axes, R.const(0, "int64"))
                lv2: R.Tensor((2,), dtype="int64") = R.add(axes, R.const(4, "int64"))
                lv3: R.Tensor((4,), dtype="int64") = R.arange(
                    R.prim_value(0), R.prim_value(4), R.prim_value(1), dtype="int64"
                )
                lv4: R.Tensor((2,), dtype="int64") = R.where(lv1, lv2, axes)
                lv5: R.Tensor((4, 1), dtype="int64") = R.expand_dims(lv3, axis=[1])
                lv6: R.Tensor((1, 2), dtype="int64") = R.expand_dims(lv4, axis=[0])
                lv7: R.Tensor((4, 2), dtype="bool") = R.equal(lv5, lv6)
                lv8: R.Tensor((4, 2), dtype="int64") = R.astype(lv7, dtype="int64")
                lv9: R.Tensor((4,), dtype="int64") = R.sum(lv8, axis=[1], keepdims=False)
                lv10: R.Tensor((4,), dtype="bool") = R.equal(lv9, R.const(0, "int64"))
                lv11: R.Tensor((1, squeeze_num_keep_dims), dtype="int64") = R.match_cast(
                    R.nonzero(lv10), R.Tensor((1, squeeze_num_keep_dims), dtype="int64")
                )
                lv12: R.Tensor((4,), dtype="int64") = R.shape_to_tensor(lv)
                lv13: R.Tensor((squeeze_num_keep_dims,), dtype="int64") = R.reshape(
                    lv11, R.shape([squeeze_num_keep_dims])
                )
                lv14: R.Tensor((2,), dtype="int64") = R.match_cast(
                    R.take(lv12, lv13, axis=0, mode="fast"), R.Tensor((2,), dtype="int64")
                )
                lv15: R.Shape(ndim=2) = R.tensor_to_shape(lv14)
                lv16: R.Shape([squeeze_dim_0, squeeze_dim_1]) = R.match_cast(
                    lv15, R.Shape([squeeze_dim_0, squeeze_dim_1])
                )
                gv: R.Tensor((squeeze_dim_0, squeeze_dim_1), dtype="float32") = R.reshape(
                    x, R.shape([squeeze_dim_0, squeeze_dim_1])
                )
                R.output(gv)
            return gv

    tvm.ir.assert_structural_equal(tvm_model, Expected)


def test_squeeze_dynamic_axes_rank_validation():
    squeeze_node = helper.make_node("Squeeze", ["x", "axes"], ["y"])
    shape = [1, 32, 1, 32]

    graph = helper.make_graph(
        [squeeze_node],
        "squeeze_dynamic_axes_rank_validation",
        inputs=[
            helper.make_tensor_value_info("x", TensorProto.FLOAT, shape),
            helper.make_tensor_value_info("axes", TensorProto.INT64, [1, 2]),
        ],
        outputs=[helper.make_tensor_value_info("y", TensorProto.FLOAT, [32, 32])],
    )

    model = helper.make_model(graph, producer_name="squeeze_dynamic_axes_rank_validation_test")
    with pytest.raises(ValueError, match="Expected a 1-D tensor"):
        from_onnx(model, opset=13, keep_params_in_input=True)


@pytest.mark.parametrize("axis", [[0]])
def test_dynamic_shape_squeeze(axis):
    shape_node = helper.make_node("Shape", ["x"], ["y"])
    squeeze_node = helper.make_node("Squeeze", ["y", "axes"], ["z"])
    shape = ["A"]

    initializer = (
        [helper.make_tensor("axes", TensorProto.INT64, [len(axis)], axis)] if axis else None
    )

    graph = helper.make_graph(
        [shape_node, squeeze_node],
        "squeeze_test",
        inputs=[
            helper.make_tensor_value_info("x", TensorProto.FLOAT, shape),
        ],
        initializer=initializer,
        outputs=[helper.make_tensor_value_info("z", TensorProto.INT64, [])],
    )

    model = helper.make_model(graph, producer_name="squeeze_test")
    tvm_model = from_onnx(model, opset=13, keep_params_in_input=True)
    assert len(tvm_model["main"].attrs["params"]) == 1
    tvm_model["main"] = tvm_model["main"].without_attr("params")

    # Use an ordinary symbolic Var for the dynamic shape binding.
    a = tvm.tirx.Var("A", "int64")
    x = relax.Var("x", relax.TensorType([a], "float32"))
    axes = relax.Var("axes", relax.TensorType([1], "int64"))
    gv = relax.Var("gv", tvm.ir.PrimType("int64"))
    body = relax.SeqExpr([relax.DataflowBlock([relax.VarBinding(gv, a)])], gv)
    # Match the importer boundary, where BlockBuilder populates the SeqExpr result type.
    body = relax.BlockBuilder().normalize(body)
    expected_func = relax.Function([x, axes], body, tvm.ir.PrimType("int64")).with_attrs(
        {"num_input": 1, "global_symbol": "main"}
    )
    tvm.ir.assert_structural_equal(tvm_model, tvm.IRModule({"main": expected_func}))


def test_const():
    shape = [32, 32]
    const_value = np.random.rand(*shape).astype(np.float32)
    const_node = helper.make_node(
        "Constant",
        [],
        ["y"],
        value=helper.make_tensor("value", TensorProto.FLOAT, shape, const_value.flatten()),
    )
    graph = helper.make_graph(
        [const_node],
        "const_test",
        inputs=[],
        outputs=[helper.make_tensor_value_info("y", TensorProto.FLOAT, shape)],
    )

    model = helper.make_model(graph, producer_name="const_test")
    tvm_model = from_onnx(model, keep_params_in_input=True)

    @I.ir_module
    class Expected:
        @R.function
        def main() -> R.Tensor((32, 32), dtype="float32"):
            R.func_attr({"num_input": 0})
            with R.dataflow():
                gv: R.Tensor((32, 32), dtype="float32") = R.const(const_value, "float32")
                R.output(gv)
            return gv

    tvm.ir.assert_structural_equal(tvm_model, Expected)


def test_instance_norm():
    def verify_instance_norm(input_shape, scale_shape, bias_shape, expected):
        node = helper.make_node("InstanceNormalization", ["a", "b", "c"], ["d"], epsilon=1e-12)
        graph = helper.make_graph(
            [node],
            "instance_norm_test",
            inputs=[
                helper.make_tensor_value_info("a", TensorProto.FLOAT, input_shape),
                helper.make_tensor_value_info("b", TensorProto.FLOAT, scale_shape),
                helper.make_tensor_value_info("c", TensorProto.FLOAT, bias_shape),
            ],
            outputs=[helper.make_tensor_value_info("d", TensorProto.FLOAT, input_shape)],
        )

        model = helper.make_model(graph, producer_name="instance_norm_test")
        tvm_model = from_onnx(model, keep_params_in_input=True)
        tvm.ir.assert_structural_equal(tvm_model, expected)

    @I.ir_module
    class Expected4D:
        @R.function
        def main(
            a: R.Tensor((1, 3, 32, 32), dtype="float32"),
            b: R.Tensor((3,), dtype="float32"),
            c: R.Tensor((3,), dtype="float32"),
        ) -> R.Tensor((1, 3, 32, 32), dtype="float32"):
            R.func_attr({"num_input": 3})
            with R.dataflow():
                lv: R.Tensor((1, 3, 1, 1), dtype="float32") = R.mean(a, axis=[2, 3], keepdims=True)
                lv1: R.Tensor((1, 3, 32, 32), dtype="float32") = R.subtract(a, lv)
                lv2: R.Tensor((1, 3, 1, 1), dtype="float32") = R.variance(
                    a, axis=[2, 3], keepdims=True
                )
                lv3: R.Tensor((1, 3, 1, 1), dtype="float32") = R.add(lv2, R.const(1e-12, "float32"))
                lv4: R.Tensor((1, 3, 1, 1), dtype="float32") = R.sqrt(lv3)
                lv5: R.Tensor((1, 3, 32, 32), dtype="float32") = R.divide(lv1, lv4)
                lv6: R.Tensor((3, 1, 1), dtype="float32") = R.reshape(b, R.shape([3, 1, 1]))
                lv7: R.Tensor((1, 3, 32, 32), dtype="float32") = R.multiply(lv5, lv6)
                lv8: R.Tensor((3, 1, 1), dtype="float32") = R.reshape(c, R.shape([3, 1, 1]))
                gv: R.Tensor((1, 3, 32, 32), dtype="float32") = R.add(lv7, lv8)
                R.output(gv)
            return gv

    @I.ir_module
    class Expected3D:
        @R.function
        def main(
            a: R.Tensor((1, 32, 32), dtype="float32"),
            b: R.Tensor((32,), dtype="float32"),
            c: R.Tensor((32,), dtype="float32"),
        ) -> R.Tensor((1, 32, 32), dtype="float32"):
            R.func_attr({"num_input": 3})
            with R.dataflow():
                lv: R.Tensor((1, 32, 1), dtype="float32") = R.mean(a, axis=[2], keepdims=True)
                lv1: R.Tensor((1, 32, 32), dtype="float32") = R.subtract(a, lv)
                lv2: R.Tensor((1, 32, 1), dtype="float32") = R.variance(a, axis=[2], keepdims=True)
                lv3: R.Tensor((1, 32, 1), dtype="float32") = R.add(lv2, R.const(1e-12, "float32"))
                lv4: R.Tensor((1, 32, 1), dtype="float32") = R.sqrt(lv3)
                lv5: R.Tensor((1, 32, 32), dtype="float32") = R.divide(lv1, lv4)
                lv6: R.Tensor((32, 1), dtype="float32") = R.reshape(b, R.shape([32, 1]))
                lv7: R.Tensor((1, 32, 32), dtype="float32") = R.multiply(lv5, lv6)
                lv8: R.Tensor((32, 1), dtype="float32") = R.reshape(c, R.shape([32, 1]))
                gv: R.Tensor((1, 32, 32), dtype="float32") = R.add(lv7, lv8)
                R.output(gv)
            return gv

    verify_instance_norm([1, 3, 32, 32], [3], [3], Expected4D)
    verify_instance_norm([1, 32, 32], [32], [32], Expected3D)


def test_mean_variance_norm():
    def verify_mean_variance_norm(axes, expected):
        node = helper.make_node("MeanVarianceNormalization", ["x"], ["y"], axes=axes)
        graph = helper.make_graph(
            [node],
            "mean_variance_norm_test",
            inputs=[helper.make_tensor_value_info("x", TensorProto.FLOAT, [1, 3, 32, 32])],
            outputs=[helper.make_tensor_value_info("y", TensorProto.FLOAT, [1, 3, 32, 32])],
        )

        model = helper.make_model(graph, producer_name="mean_variance_norm_test")
        tvm_model = from_onnx(model, keep_params_in_input=True)
        tvm.ir.assert_structural_equal(tvm_model, expected)

    @I.ir_module
    class ExpectedDefaultAxes:
        @R.function
        def main(
            x: R.Tensor((1, 3, 32, 32), dtype="float32"),
        ) -> R.Tensor((1, 3, 32, 32), dtype="float32"):
            R.func_attr({"num_input": 1})
            with R.dataflow():
                lv: R.Tensor((1, 3, 1, 1), dtype="float32") = R.mean(
                    x, axis=[0, 2, 3], keepdims=True
                )
                lv1: R.Tensor((1, 3, 32, 32), dtype="float32") = R.subtract(x, lv)
                lv2: R.Tensor((1, 3, 32, 32), dtype="float32") = R.power(x, R.const(2.0, "float32"))
                lv3: R.Tensor((1, 3, 1, 1), dtype="float32") = R.mean(
                    lv2, axis=[0, 2, 3], keepdims=True
                )
                lv4: R.Tensor((1, 3, 1, 1), dtype="float32") = R.power(lv, R.const(2.0, "float32"))
                lv5: R.Tensor((1, 3, 1, 1), dtype="float32") = R.subtract(lv3, lv4)
                lv6: R.Tensor((1, 3, 1, 1), dtype="float32") = R.sqrt(lv5)
                gv: R.Tensor((1, 3, 32, 32), dtype="float32") = R.divide(lv1, lv6)
                R.output(gv)
            return gv

    @I.ir_module
    class ExpectedChannelAxes:
        @R.function
        def main(
            x: R.Tensor((1, 3, 32, 32), dtype="float32"),
        ) -> R.Tensor((1, 3, 32, 32), dtype="float32"):
            R.func_attr({"num_input": 1})
            with R.dataflow():
                lv: R.Tensor((1, 1, 1, 1), dtype="float32") = R.mean(
                    x, axis=[1, 2, 3], keepdims=True
                )
                lv1: R.Tensor((1, 3, 32, 32), dtype="float32") = R.subtract(x, lv)
                lv2: R.Tensor((1, 3, 32, 32), dtype="float32") = R.power(x, R.const(2.0, "float32"))
                lv3: R.Tensor((1, 1, 1, 1), dtype="float32") = R.mean(
                    lv2, axis=[1, 2, 3], keepdims=True
                )
                lv4: R.Tensor((1, 1, 1, 1), dtype="float32") = R.power(lv, R.const(2.0, "float32"))
                lv5: R.Tensor((1, 1, 1, 1), dtype="float32") = R.subtract(lv3, lv4)
                lv6: R.Tensor((1, 1, 1, 1), dtype="float32") = R.sqrt(lv5)
                gv: R.Tensor((1, 3, 32, 32), dtype="float32") = R.divide(lv1, lv6)
                R.output(gv)
            return gv

    verify_mean_variance_norm((0, 2, 3), ExpectedDefaultAxes)
    verify_mean_variance_norm((1, 2, 3), ExpectedChannelAxes)


def test_layer_norm():
    layer_norm_node = helper.make_node(
        "LayerNormalization", ["input", "scale", "bias"], ["Y"], epsilon=1e-12
    )

    graph = helper.make_graph(
        [layer_norm_node],
        "layer_norm_test",
        inputs=[
            helper.make_tensor_value_info("input", TensorProto.FLOAT, [32, 32]),
            helper.make_tensor_value_info("scale", TensorProto.FLOAT, [32]),
            helper.make_tensor_value_info("bias", TensorProto.FLOAT, [32]),
        ],
        outputs=[
            helper.make_tensor_value_info("Y", TensorProto.FLOAT, [32, 32]),
        ],
    )

    model = helper.make_model(graph, producer_name="layer_norm_test")
    check_correctness(model)

    # Test case with no bias that is an optional input
    layer_norm_node = helper.make_node(
        "LayerNormalization", ["input", "scale"], ["Y"], epsilon=1e-12
    )

    graph = helper.make_graph(
        [layer_norm_node],
        "layer_norm_test",
        inputs=[
            helper.make_tensor_value_info("input", TensorProto.FLOAT, [32, 32]),
            helper.make_tensor_value_info("scale", TensorProto.FLOAT, [32]),
        ],
        outputs=[
            helper.make_tensor_value_info("Y", TensorProto.FLOAT, [32, 32]),
        ],
    )

    model = helper.make_model(graph, producer_name="layer_norm_test")
    check_correctness(model)

    # No bias with a non-square input where data.shape[1] differs from the scale
    # shape, see https://github.com/apache/tvm/issues/19691.
    layer_norm_node = helper.make_node(
        "LayerNormalization", ["input", "scale"], ["Y"], axis=-1, epsilon=1e-12
    )

    graph = helper.make_graph(
        [layer_norm_node],
        "layer_norm_test",
        inputs=[
            helper.make_tensor_value_info("input", TensorProto.FLOAT, [2, 3, 4, 8]),
            helper.make_tensor_value_info("scale", TensorProto.FLOAT, [8]),
        ],
        outputs=[
            helper.make_tensor_value_info("Y", TensorProto.FLOAT, [2, 3, 4, 8]),
        ],
    )

    model = helper.make_model(graph, producer_name="layer_norm_test")
    check_correctness(model)

    # No bias with a non-square fp16 input. The synthesized zero bias must match
    # the scale dtype, otherwise layer_norm rejects the float32 bias, see
    # https://github.com/apache/tvm/issues/19691.
    layer_norm_node = helper.make_node(
        "LayerNormalization", ["input", "scale"], ["Y"], axis=-1, epsilon=1e-12
    )

    graph = helper.make_graph(
        [layer_norm_node],
        "layer_norm_test",
        inputs=[
            helper.make_tensor_value_info("input", TensorProto.FLOAT16, [2, 3, 4, 8]),
            helper.make_tensor_value_info("scale", TensorProto.FLOAT16, [8]),
        ],
        outputs=[
            helper.make_tensor_value_info("Y", TensorProto.FLOAT16, [2, 3, 4, 8]),
        ],
    )

    model = helper.make_model(graph, producer_name="layer_norm_test")
    check_correctness(model, opset=17, atol=1e-2, rtol=1e-2)

    # Same no-bias path for bf16. ONNX Runtime's CPU provider has no bf16
    # LayerNormalization kernel, so this only checks the importer builds the
    # graph with a bf16 zero bias (the dtype the fix derives from the scale).
    layer_norm_node = helper.make_node(
        "LayerNormalization", ["input", "scale"], ["Y"], axis=-1, epsilon=1e-12
    )

    graph = helper.make_graph(
        [layer_norm_node],
        "layer_norm_test",
        inputs=[
            helper.make_tensor_value_info("input", TensorProto.BFLOAT16, [2, 3, 4, 8]),
            helper.make_tensor_value_info("scale", TensorProto.BFLOAT16, [8]),
        ],
        outputs=[
            helper.make_tensor_value_info("Y", TensorProto.BFLOAT16, [2, 3, 4, 8]),
        ],
    )

    model = helper.make_model(graph, producer_name="layer_norm_test")
    model.opset_import[0].version = 17
    from_onnx(model, opset=17, keep_params_in_input=True)


def test_layer_norm_with_nd_gamma_beta():
    layer_norm_node = helper.make_node(
        "LayerNormalization", ["input", "scale", "bias"], ["Y"], axis=1, epsilon=1e-12
    )

    graph = helper.make_graph(
        [layer_norm_node],
        "layer_norm_with_nd_gamma_beta_test",
        inputs=[
            helper.make_tensor_value_info("input", TensorProto.FLOAT, [1, 3, 4, 4]),
            helper.make_tensor_value_info("scale", TensorProto.FLOAT, [3, 4, 4]),
            helper.make_tensor_value_info("bias", TensorProto.FLOAT, [3, 4, 4]),
        ],
        outputs=[
            helper.make_tensor_value_info("Y", TensorProto.FLOAT, [1, 3, 4, 4]),
        ],
    )

    model = helper.make_model(graph, producer_name="layer_norm_with_nd_gamma_beta_test")
    check_correctness(model)

    # Test case with no bias that is an optional input
    layer_norm_node = helper.make_node(
        "LayerNormalization", ["input", "scale"], ["Y"], axis=1, epsilon=1e-12
    )

    graph = helper.make_graph(
        [layer_norm_node],
        "layer_norm_with_nd_gamma_beta_test",
        inputs=[
            helper.make_tensor_value_info("input", TensorProto.FLOAT, [32, 32]),
            helper.make_tensor_value_info("scale", TensorProto.FLOAT, [32]),
        ],
        outputs=[
            helper.make_tensor_value_info("Y", TensorProto.FLOAT, [32, 32]),
        ],
    )

    model = helper.make_model(graph, producer_name="layer_norm_with_nd_gamma_beta_test")
    check_correctness(model)


def test_layer_norm_numerical_stability():
    """Numerical stability test for https://github.com/apache/tvm/issues/19592."""
    layer_norm_node = helper.make_node(
        "LayerNormalization", ["input", "scale", "bias"], ["Y"], axis=-1, epsilon=1e-5
    )
    graph = helper.make_graph(
        [layer_norm_node],
        "layer_norm_numerical_stability",
        inputs=[
            helper.make_tensor_value_info("input", TensorProto.FLOAT, [1, 4]),
            helper.make_tensor_value_info("scale", TensorProto.FLOAT, [4]),
            helper.make_tensor_value_info("bias", TensorProto.FLOAT, [4]),
        ],
        outputs=[
            helper.make_tensor_value_info("Y", TensorProto.FLOAT, [1, 4]),
        ],
    )
    model = helper.make_model(graph, producer_name="layer_norm_numerical_stability")

    input_array = np.array([[80000.0, 80001.0, 80002.0, 80003.0]], dtype=np.float32)
    scale_array = np.ones(4, dtype=np.float32)
    bias_array = np.zeros(4, dtype=np.float32)
    inputs = {"input": input_array, "scale": scale_array, "bias": bias_array}

    # ONNXRuntime also returns NaN for Large-value, small-variance inputs, so we here
    # compare against a two-pass reference instead of ORT.
    mean = input_array.mean(axis=-1, keepdims=True)
    var = ((input_array - mean) ** 2).mean(axis=-1, keepdims=True)
    expected = ((input_array - mean) / np.sqrt(var + 1e-5) * scale_array + bias_array).astype(
        np.float32
    )

    tvm_output = run_in_tvm(model, inputs=inputs, ir_version=9, opset=17)

    assert np.isfinite(tvm_output.numpy()).all()
    tvm.testing.assert_allclose(tvm_output.numpy(), expected)


def test_rms_norm():
    # Basic test: default axis=-1
    rms_norm_node = helper.make_node("RMSNormalization", ["input", "scale"], ["Y"], epsilon=1e-05)

    graph = helper.make_graph(
        [rms_norm_node],
        "rms_norm_test",
        inputs=[
            helper.make_tensor_value_info("input", TensorProto.FLOAT, [2, 8, 32]),
            helper.make_tensor_value_info("scale", TensorProto.FLOAT, [32]),
        ],
        outputs=[
            helper.make_tensor_value_info("Y", TensorProto.FLOAT, [2, 8, 32]),
        ],
    )

    model = helper.make_model(graph, producer_name="rms_norm_test")
    check_correctness(model, opset=23)

    # Test with explicit axis=1 (normalize over last 2 dims)
    rms_norm_node = helper.make_node(
        "RMSNormalization", ["input", "scale"], ["Y"], axis=1, epsilon=1e-06
    )

    graph = helper.make_graph(
        [rms_norm_node],
        "rms_norm_axis_test",
        inputs=[
            helper.make_tensor_value_info("input", TensorProto.FLOAT, [4, 8, 16]),
            helper.make_tensor_value_info("scale", TensorProto.FLOAT, [8, 16]),
        ],
        outputs=[
            helper.make_tensor_value_info("Y", TensorProto.FLOAT, [4, 8, 16]),
        ],
    )

    model = helper.make_model(graph, producer_name="rms_norm_axis_test")
    check_correctness(model, opset=23)

    # Test with float16 input (stash_type=1 means compute in float32)
    rms_norm_node = helper.make_node(
        "RMSNormalization", ["input", "scale"], ["Y"], epsilon=1e-05, stash_type=1
    )

    graph = helper.make_graph(
        [rms_norm_node],
        "rms_norm_fp16_test",
        inputs=[
            helper.make_tensor_value_info("input", TensorProto.FLOAT16, [2, 8, 32]),
            helper.make_tensor_value_info("scale", TensorProto.FLOAT16, [32]),
        ],
        outputs=[
            helper.make_tensor_value_info("Y", TensorProto.FLOAT16, [2, 8, 32]),
        ],
    )

    model = helper.make_model(graph, producer_name="rms_norm_fp16_test")
    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):
    def verify_skiplayernormalization(input_, skip, gamma, beta, bias):
        node = onnx.helper.make_node(
            "SkipLayerNormalization",
            inputs=["input", "skip", "gamma", "beta", "bias"],
            outputs=["output", "mean", "std_dev"],
            domain="com.microsoft",
        )

        node.attribute.append(onnx.helper.make_attribute("epsilon", 1e-4))

        input_shape = list(input_.shape)
        skip_shape = list(skip.shape)
        gamma_shape = list(gamma.shape)
        beta_shape = list(beta.shape)
        bias_shape = list(bias.shape)
        output_shape = list(input_.shape)
        mean_shape = list([1])
        std_dev_shape = list([1])
        if dynamic:
            input_shape = ["?" for _ in range(len(input_.shape))]
            skip_shape = ["?" for _ in range(len(skip.shape))]
            gamma_shape = ["?" for _ in range(len(gamma.shape))]
            beta_shape = ["?" for _ in range(len(beta.shape))]
            bias_shape = ["?" for _ in range(len(bias.shape))]
            output_shape = ["?" for _ in range(len(input_.shape))]

        graph = helper.make_graph(
            [node],
            "skiplayernormalization_test",
            inputs=[
                helper.make_tensor_value_info("input", TensorProto.FLOAT, input_shape),
                helper.make_tensor_value_info("skip", TensorProto.FLOAT, skip_shape),
                helper.make_tensor_value_info("gamma", TensorProto.FLOAT, gamma_shape),
                helper.make_tensor_value_info("beta", TensorProto.FLOAT, beta_shape),
                helper.make_tensor_value_info("bias", TensorProto.FLOAT, bias_shape),
            ],
            outputs=[
                helper.make_tensor_value_info("output", TensorProto.FLOAT, output_shape),
                helper.make_tensor_value_info("mean", TensorProto.FLOAT, mean_shape),
                helper.make_tensor_value_info("std_dev", TensorProto.FLOAT, std_dev_shape),
            ],
        )

        model = helper.make_model(graph, producer_name="skiplayernormalization_test")
        tvm_model = from_onnx(model, keep_params_in_input=True)

        @I.ir_module
        class Expected:
            @R.function
            def main(
                input: R.Tensor((4, 4, 384), dtype="float32"),
                skip: R.Tensor((4, 4, 384), dtype="float32"),
                gamma: R.Tensor((384,), dtype="float32"),
                beta: R.Tensor((384,), dtype="float32"),
                bias: R.Tensor((384,), dtype="float32"),
            ) -> R.Tuple(
                R.Tensor((4, 4, 384), dtype="float32"),
                R.Tensor((), dtype="float32"),
                R.Tensor((), dtype="float32"),
            ):
                R.func_attr({"num_input": 5})
                with R.dataflow():
                    lv: R.Tensor((4, 4, 384), dtype="float32") = R.add(input, skip)
                    lv1: R.Tensor((4, 4, 384), dtype="float32") = R.add(lv, bias)
                    lv2: R.Tensor((4, 4, 384), dtype="float32") = R.nn.layer_norm(
                        lv1, gamma, beta, axes=-1, epsilon=9.999999747378752e-05
                    )
                    gv: R.Tuple(
                        R.Tensor((4, 4, 384), dtype="float32"),
                        R.Tensor((), dtype="float32"),
                        R.Tensor((), dtype="float32"),
                    ) = (lv2, R.const(0, "float32"), R.const(0, "float32"))
                    R.output(gv)
                return gv

        tvm.ir.assert_structural_equal(tvm_model, Expected)

    hidden_size = 384
    batch_size = 4
    sequence_length = 4

    dtype = "float32"
    input_array = np.random.random((batch_size, sequence_length, hidden_size)).astype(dtype)
    skip = np.random.random((batch_size, sequence_length, hidden_size)).astype(dtype)
    gamma = np.random.uniform(0.5, 0.7, hidden_size).astype(dtype)
    beta = np.random.randn(hidden_size).astype(dtype) * 0.1
    bias = np.random.randn(hidden_size).astype(dtype)

    verify_skiplayernormalization(input_array, skip, gamma, beta, bias)


def test_embedlayernormalization():
    def verify_embedlayernormalization(
        input_ids,
        segment_ids,
        word_embedding,
        position_embedding,
        segment_embedding,
        gamma,
        beta,
        expected,
    ):
        node = onnx.helper.make_node(
            "EmbedLayerNormalization",
            inputs=[
                "input_ids",
                "" if segment_ids is None else "segment_ids",
                "word_embedding",
                "position_embedding",
                "" if segment_embedding is None else "segment_embedding",
                "gamma",
                "beta",
            ],
            outputs=["output", "mask_index"],
            domain="com.microsoft",
        )

        node.attribute.append(onnx.helper.make_attribute("epsilon", 1e-4))

        segment_ids_shape = [] if segment_ids is None else segment_ids.shape
        segment_embedding_shape = [] if segment_embedding is None else segment_embedding.shape

        graph = helper.make_graph(
            [node],
            "embedlayernormalization_test",
            inputs=[
                helper.make_tensor_value_info(
                    "input_ids", TensorProto.INT32, list(input_ids.shape)
                ),
                helper.make_tensor_value_info("segment_ids", TensorProto.INT32, segment_ids_shape),
                helper.make_tensor_value_info(
                    "word_embedding", TensorProto.FLOAT, list(word_embedding.shape)
                ),
                helper.make_tensor_value_info(
                    "position_embedding", TensorProto.FLOAT, list(position_embedding.shape)
                ),
                helper.make_tensor_value_info(
                    "segment_embedding", TensorProto.FLOAT, segment_embedding_shape
                ),
                helper.make_tensor_value_info("gamma", TensorProto.FLOAT, list(gamma.shape)),
                helper.make_tensor_value_info("beta", TensorProto.FLOAT, list(beta.shape)),
            ],
            outputs=[
                helper.make_tensor_value_info(
                    "output", TensorProto.FLOAT, list((batch_size, sequence_length, hidden_size))
                ),
                helper.make_tensor_value_info("mask_index", TensorProto.INT32, [batch_size]),
            ],
        )

        model = helper.make_model(graph, producer_name="embedlayernormalization_test")

        tvm_model = from_onnx(model, keep_params_in_input=True)

        tvm.ir.assert_structural_equal(tvm_model, expected)

        # TODO(@anwang2009): onnxruntime v1.9.0 requires empty list for optional argument,
        # but v1.10.0+ requires None instead.
        # verify_with_ort_with_inputs(
        #     model,
        #     [
        #         input_ids,
        #         np.empty(0, dtype="int32") if segment_ids is None else segment_ids,
        #         word_embedding,
        #         position_embedding,
        #         np.empty(0, dtype="float32") if segment_embedding is None else segment_embedding,
        #         gamma,
        #         beta,
        #     ],
        #     [
        #         (batch_size, sequence_length, hidden_size),
        #         batch_size,
        #     ],
        #     target=target,
        #     dev=dev,
        #     rtol=1e-4,
        #     atol=1e-4,
        # )

    hidden_size = 384
    batch_size = 4
    sequence_length = 3
    vocab_size = 5

    input_ids = np.full((batch_size, sequence_length), 3).astype("int32")
    segment_ids = np.zeros((batch_size, sequence_length)).astype("int32")
    word_embedding = np.full((vocab_size, hidden_size), 1).astype("float32")
    position_embedding = np.full((sequence_length, hidden_size), 2).astype("float32")
    segment_embedding = np.full((vocab_size, hidden_size), 3).astype("float32")

    gamma = np.random.uniform(0.5, 0.7, hidden_size).astype("float32")
    beta = np.random.randn(hidden_size).astype("float32") * 0.1

    @I.ir_module
    class ExpectedNoSegment:
        @R.function
        def main(
            input_ids: R.Tensor((4, 3), dtype="int32"),
            segment_ids: R.Tensor((), dtype="int32"),
            word_embedding: R.Tensor((5, 384), dtype="float32"),
            position_embedding: R.Tensor((3, 384), dtype="float32"),
            segment_embedding: R.Tensor((), dtype="float32"),
            gamma: R.Tensor((384,), dtype="float32"),
            beta: R.Tensor((384,), dtype="float32"),
        ) -> R.Tuple(
            R.Tensor((4, 3, 384), dtype="float32"),
            R.Tensor((4,), dtype="int32"),
        ):
            R.func_attr({"num_input": 7})
            with R.dataflow():
                lv: R.Tensor((4, 3, 384), dtype="float32") = R.take(
                    word_embedding, input_ids, axis=0, mode="fast"
                )
                lv1: R.Tensor((4, 3, 384), dtype="float32") = R.take(
                    position_embedding,
                    R.const([[0, 1, 2], [0, 1, 2], [0, 1, 2], [0, 1, 2]], "int64"),
                    axis=0,
                    mode="fast",
                )
                lv2: R.Tensor((4, 3, 384), dtype="float32") = R.add(lv, lv1)
                lv3: R.Tensor((4, 3, 384), dtype="float32") = R.nn.layer_norm(
                    lv2, gamma, beta, axes=-1, epsilon=9.999999747378752e-05
                )
                gv: R.Tuple(
                    R.Tensor((4, 3, 384), dtype="float32"),
                    R.Tensor((4,), dtype="int32"),
                ) = (lv3, R.const([0, 0, 0, 0], "int32"))
                R.output(gv)
            return gv

    @I.ir_module
    class ExpectedWithSegment:
        @R.function
        def main(
            input_ids: R.Tensor((4, 3), dtype="int32"),
            segment_ids: R.Tensor((4, 3), dtype="int32"),
            word_embedding: R.Tensor((5, 384), dtype="float32"),
            position_embedding: R.Tensor((3, 384), dtype="float32"),
            segment_embedding: R.Tensor((5, 384), dtype="float32"),
            gamma: R.Tensor((384,), dtype="float32"),
            beta: R.Tensor((384,), dtype="float32"),
        ) -> R.Tuple(
            R.Tensor((4, 3, 384), dtype="float32"),
            R.Tensor((4,), dtype="int32"),
        ):
            R.func_attr({"num_input": 7})
            with R.dataflow():
                lv: R.Tensor((4, 3, 384), dtype="float32") = R.take(
                    word_embedding, input_ids, axis=0, mode="fast"
                )
                lv1: R.Tensor((4, 3, 384), dtype="float32") = R.take(
                    position_embedding,
                    R.const([[0, 1, 2], [0, 1, 2], [0, 1, 2], [0, 1, 2]], "int64"),
                    axis=0,
                    mode="fast",
                )
                lv2: R.Tensor((4, 3, 384), dtype="float32") = R.add(lv, lv1)
                lv3: R.Tensor((4, 3, 384), dtype="float32") = R.take(
                    segment_embedding, segment_ids, axis=0, mode="fast"
                )
                lv4: R.Tensor((4, 3, 384), dtype="float32") = R.add(lv2, lv3)
                lv5: R.Tensor((4, 3, 384), dtype="float32") = R.nn.layer_norm(
                    lv4, gamma, beta, axes=-1, epsilon=9.999999747378752e-05
                )
                gv: R.Tuple(
                    R.Tensor((4, 3, 384), dtype="float32"),
                    R.Tensor((4,), dtype="int32"),
                ) = (lv5, R.const([0, 0, 0, 0], "int32"))
                R.output(gv)
            return gv

    verify_embedlayernormalization(
        input_ids,
        segment_ids,
        word_embedding,
        position_embedding,
        segment_embedding,
        gamma,
        beta,
        ExpectedWithSegment,
    )

    # Test with undefined segment embedding
    verify_embedlayernormalization(
        input_ids,
        None,
        word_embedding,
        position_embedding,
        None,
        gamma,
        beta,
        ExpectedNoSegment,
    )


def test_local_response_norm():
    lrn_node = helper.make_node(
        op_type="LRN",
        inputs=["input"],
        outputs=["output"],
        name="LRN_Node",
        alpha=0.0001,
        beta=0.75,
        bias=1.0,
        size=3,
    )

    graph = helper.make_graph(
        [lrn_node],
        "local_response_norm_test",
        inputs=[
            helper.make_tensor_value_info("input", TensorProto.FLOAT, [1, 3, 32, 32]),
        ],
        outputs=[
            helper.make_tensor_value_info("output", TensorProto.FLOAT, [1, 3, 32, 32]),
        ],
    )

    model = helper.make_model(graph, producer_name="local_response_norm_test")
    tvm_model = from_onnx(model, keep_params_in_input=True)

    @I.ir_module
    class Expected:
        @R.function
        def main(
            input: R.Tensor((1, 3, 32, 32), dtype="float32"),
        ) -> R.Tensor((1, 3, 32, 32), dtype="float32"):
            R.func_attr({"num_input": 1})
            with R.dataflow():
                lv: R.Tensor((1, 3, 32, 32), dtype="float32") = R.multiply(input, input)
                lv1: R.Tensor((1, 1, 3, 32, 32), dtype="float32") = R.expand_dims(lv, axis=[1])
                lv2: R.Tensor((1, 1, 3, 32, 32), dtype="float32") = R.nn.avg_pool3d(
                    lv1,
                    pool_size=[3, 1, 1],
                    strides=[1, 1, 1],
                    dilation=[1, 1, 1],
                    padding=[1, 0, 0, 1, 0, 0],
                    ceil_mode=False,
                    count_include_pad=True,
                    layout="NCDHW",
                    out_layout="NCDHW",
                )
                lv3: R.Tensor((1, 3, 32, 32), dtype="float32") = R.squeeze(lv2, axis=[1])
                lv4: R.Tensor((1, 3, 32, 32), dtype="float32") = R.multiply(
                    lv3, R.const(9.9999997473787516e-05, "float32")
                )
                lv5: R.Tensor((1, 3, 32, 32), dtype="float32") = R.add(lv4, R.const(1.0, "float32"))
                lv6: R.Tensor((1, 3, 32, 32), dtype="float32") = R.power(
                    lv5, R.const(0.75, "float32")
                )
                gv: R.Tensor((1, 3, 32, 32), dtype="float32") = R.divide(input, lv6)
                R.output(gv)
            return gv

    tvm.ir.assert_structural_equal(tvm_model, Expected)


COMPOSITE_REDUCE_FUNCS = [
    "ReduceSumSquare",
    "ReduceLogSum",
    "ReduceLogSumExp",
    "ReduceL1",
    "ReduceL2",
]

REDUCE_AXES_ATTR_TEST_CASES = [
    ([3, 2, 2], None),
    ([3, 2, 3], None),
    ([3, 3, 3], (1,)),
    ([3, 3, 3, 1], (1, 2)),
    ([3, 3, 3, 1], (1,)),
    ([1, 3, 4, 1], (1,)),
]

REDUCE_AXES_INPUT_TEST_CASES = [
    ([3, 2, 2], [], False),
    ([3, 2, 2], None, False),
    ([4, 3], [], True),
    ([3, 3, 3, 1], (1, 2), False),
]


def _reduce_output_shape(input_shape: list[int], axes, keepdims: bool, noop_with_empty_axes=False):
    if noop_with_empty_axes and not axes:
        return list(input_shape)
    axis = None if not axes else axes
    return list(np.sum(np.empty(input_shape), axis=axis, keepdims=keepdims).shape)


def verify_composite_reduce_axes_attr_ir(
    func: str,
    input_shape: list[int],
    axes,
    keepdims: bool,
    dynamic: bool,
    opset: int,
    expected,
):
    attrs = {"keepdims": keepdims}
    if axes:
        attrs["axes"] = axes
    node = onnx.helper.make_node(func, inputs=["x"], outputs=["y"], **attrs)
    output_shape = _reduce_output_shape(input_shape, axes, keepdims)
    graph = helper.make_graph(
        [node],
        "composite_reduce_axes_attr_ir_test",
        inputs=[
            helper.make_tensor_value_info(
                "x", TensorProto.FLOAT, ["?"] * len(input_shape) if dynamic else input_shape
            )
        ],
        outputs=[
            helper.make_tensor_value_info(
                "y", TensorProto.FLOAT, ["?"] * len(output_shape) if dynamic else output_shape
            )
        ],
    )
    model = helper.make_model(
        graph,
        producer_name="composite_reduce_axes_attr_ir_test",
        opset_imports=[helper.make_opsetid("", opset)],
    )
    tvm_model = from_onnx(model, opset=opset, keep_params_in_input=True)
    tvm.ir.assert_structural_equal(tvm_model, expected, map_free_vars=dynamic)


def create_reduce_test_parameters_axes_attr():
    output = []
    for value in [True, False]:
        output.append(("ReduceMax", value, 11))
        output.append(("ReduceMean", value, 13))
        output.append(("ReduceMin", value, 11))
        output.append(("ReduceProd", value, 13))
        output.append(("ReduceSum", value, 11))
        # Opset 11-12 axes-as-attr: verifies get_converter does not
        # underflow to the v18 (axes-as-input) implementation.
        output.append(("ReduceMean", value, 11))
        output.append(("ReduceProd", value, 11))
    return output


def create_composite_reduce_test_parameters_axes_attr():
    output = []
    for dynamic in [True, False]:
        for opset in [13, 11]:
            for func in COMPOSITE_REDUCE_FUNCS:
                output.append((func, dynamic, opset))
    return output


@pytest.mark.parametrize("func, dynamic, opset", create_reduce_test_parameters_axes_attr())
def test_all_reduce_funcs_axes_attr(func, dynamic, opset):
    def verify_reduce_func(func, data, axis, keepdims):
        inshape = data.shape
        outshape = np.sum(data, axis=axis, keepdims=keepdims == 1).shape

        if axis:
            node = onnx.helper.make_node(
                func, inputs=["x"], outputs=["y"], axes=axis, keepdims=keepdims
            )
        else:
            node = onnx.helper.make_node(func, inputs=["x"], outputs=["y"], keepdims=keepdims)

        if dynamic:
            in_list = ["?" for _ in range(len(inshape))]
            out_list = ["?" for _ in range(len(outshape))]
        else:
            in_list = list(inshape)
            out_list = list(outshape)
        graph = helper.make_graph(
            [node],
            "reduce_test",
            inputs=[helper.make_tensor_value_info("x", TensorProto.FLOAT, in_list)],
            outputs=[helper.make_tensor_value_info("y", TensorProto.FLOAT, out_list)],
        )

        model = helper.make_model(graph, producer_name="reduce_test")
        inputs_dict = {"x": data}
        # Reduction ops accumulate arithmetic errors, so we use a higher tolerance.
        check_correctness(model, inputs_dict, opset=opset, rtol=1e-4, atol=1e-4)

    for keepdims in [True, False]:
        verify_reduce_func(
            func, np.random.randn(3, 2, 2).astype(np.float32), axis=None, keepdims=keepdims
        )

        verify_reduce_func(
            func, np.random.randn(3, 2, 3).astype(np.float32), axis=None, keepdims=keepdims
        )

        verify_reduce_func(
            func, np.random.randn(3, 3, 3).astype(np.float32), axis=(1,), keepdims=keepdims
        )

        verify_reduce_func(
            func, np.random.randn(3, 3, 3, 1).astype(np.float32), axis=(1, 2), keepdims=keepdims
        )

        verify_reduce_func(
            func, np.random.randn(3, 3, 3, 1).astype(np.float32), axis=(1,), keepdims=keepdims
        )

        verify_reduce_func(
            func, np.random.randn(1, 3, 4, 1).astype(np.float32), axis=(1,), keepdims=keepdims
        )


def _make_composite_reduce_expected_ir(
    func: str,
    input_shape: list[int],
    axes,
    noop_with_empty_axes: bool,
    keepdims: bool,
    dynamic: bool,
    axes_as_input: bool = False,
):
    def expected_input_shape(shape):
        if not dynamic:
            return tuple(shape)
        return tuple(f"reduce_dim_{i}" for i in range(len(shape)))

    axis = None if not axes else tuple(axes)
    parser_vars = {
        "I": I,
        "R": R,
        "input_shape": expected_input_shape(input_shape),
        "axis": axis,
        "keepdims": keepdims,
    }
    params = ['        x: R.Tensor(input_shape, dtype="float32")']
    if axes_as_input and axes is not None:
        axes_shape = tuple(np.asarray(axes, dtype=np.int64).shape)
        parser_vars["axes_shape"] = axes_shape
        params.append('        reduce_axes: R.Tensor(axes_shape, dtype="int64")')

    if noop_with_empty_axes and not axes:
        body = ["            gv = x"]
    elif func == "ReduceSumSquare":
        body = [
            "            lv = R.multiply(x, x)",
            "            gv = R.sum(lv, axis=axis, keepdims=keepdims)",
        ]
    elif func == "ReduceLogSum":
        body = [
            "            lv = R.sum(x, axis=axis, keepdims=keepdims)",
            "            gv = R.log(lv)",
        ]
    elif func == "ReduceLogSumExp":
        parser_vars["logsumexp_keepdims"] = True
        body = [
            "            lv = R.max(x, axis=axis, keepdims=logsumexp_keepdims)",
            "            lv1 = R.subtract(x, lv)",
            "            lv2 = R.exp(lv1)",
            "            lv3 = R.sum(lv2, axis=axis, keepdims=logsumexp_keepdims)",
            "            lv4 = R.log(lv3)",
        ]
        if keepdims:
            body.append("            gv = R.add(lv4, lv)")
        else:
            parser_vars["squeeze_axis"] = None if axis is None else list(axis)
            body += [
                "            lv5 = R.add(lv4, lv)",
                "            gv = R.squeeze(lv5, axis=squeeze_axis)",
            ]
    elif func == "ReduceL1":
        body = [
            "            lv = R.abs(x)",
            "            gv = R.sum(lv, axis=axis, keepdims=keepdims)",
        ]
    elif func == "ReduceL2":
        body = [
            "            lv = R.multiply(x, x)",
            "            lv1 = R.sum(lv, axis=axis, keepdims=keepdims)",
            "            gv = R.sqrt(lv1)",
        ]
    else:
        raise AssertionError(f"No composite reduce expected IR for {func}")

    source = "\n".join(
        [
            "@I.ir_module",
            "class Expected:",
            "    @R.function",
            "    def main(",
            ",\n".join(params),
            "    ):",
            '        R.func_attr({"num_input": 1})',
            "        with R.dataflow():",
            *body,
            "            R.output(gv)",
            "        return gv",
            "",
        ]
    )
    return tvm.script.from_source(source, extra_vars=parser_vars)


def test_composite_reduce_funcs_axes_attr_ir():
    for func in COMPOSITE_REDUCE_FUNCS:
        for keepdims in [True, False]:
            for dynamic in [True, False]:
                for input_shape, axes in REDUCE_AXES_ATTR_TEST_CASES:
                    expected = _make_composite_reduce_expected_ir(
                        func, input_shape, axes, False, keepdims, dynamic
                    )
                    for opset in [13, 11]:
                        verify_composite_reduce_axes_attr_ir(
                            func, input_shape, axes, keepdims, dynamic, opset, expected
                        )


def create_reduce_test_parameters_axes_input():
    output = []
    for dynamic in [True, False]:
        output.append(("ReduceMax", dynamic, 18))
        output.append(("ReduceMean", dynamic, 18))
        output.append(("ReduceMin", dynamic, 18))
        output.append(("ReduceProd", dynamic, 18))
        output.append(("ReduceSum", dynamic, 13))
    return output


def create_composite_reduce_test_parameters_axes_input():
    output = []
    for dynamic in [True, False]:
        for func in COMPOSITE_REDUCE_FUNCS:
            output.append((func, dynamic, 18))
    return output


def verify_composite_reduce_axes_input_ir(
    func: str,
    input_shape: list[int],
    axes,
    noop_with_empty_axes: bool,
    keepdims: bool,
    dynamic: bool,
    opset: int,
    expected,
):
    node_inputs = ["x"]
    initializers = []
    axes_input_shape = None
    if axes is not None:
        axes_np = np.asarray(axes, dtype=np.int64)
        axes_input_shape = list(axes_np.shape)
        initializers.append(
            helper.make_tensor(
                name="reduce_axes",
                data_type=TensorProto.INT64,
                dims=axes_input_shape,
                vals=axes_np,
            )
        )
        node_inputs.append("reduce_axes")

    effective_axes = None if not axes and not noop_with_empty_axes else axes
    output_shape = _reduce_output_shape(
        input_shape, effective_axes, keepdims, noop_with_empty_axes=noop_with_empty_axes
    )
    node = onnx.helper.make_node(
        func,
        inputs=node_inputs,
        outputs=["y"],
        keepdims=keepdims,
        noop_with_empty_axes=noop_with_empty_axes,
    )
    graph = helper.make_graph(
        [node],
        "composite_reduce_axes_input_ir_test",
        inputs=[
            helper.make_tensor_value_info(
                "x", TensorProto.FLOAT, ["?"] * len(input_shape) if dynamic else input_shape
            )
        ],
        initializer=initializers,
        outputs=[
            helper.make_tensor_value_info(
                "y", TensorProto.FLOAT, ["?"] * len(output_shape) if dynamic else output_shape
            )
        ],
    )
    model = helper.make_model(
        graph,
        producer_name="composite_reduce_axes_input_ir_test",
        opset_imports=[helper.make_opsetid("", opset)],
    )
    tvm_model = from_onnx(model, opset=opset, keep_params_in_input=True)
    if axes_input_shape is not None:
        assert len(tvm_model["main"].attrs["params"]) == 1
        tvm_model["main"] = tvm_model["main"].without_attr("params")

    tvm.ir.assert_structural_equal(tvm_model, expected, map_free_vars=dynamic)


@pytest.mark.parametrize("func, dynamic, opset", create_reduce_test_parameters_axes_input())
def test_all_reduce_funcs_axes_input(func, dynamic, opset):
    def verify_reduce_func(func, data, axes, keepdims, noop_with_empty_axes=False):
        inshape = data.shape
        inputs = ["x"]
        initializers = []

        # Optional `axes` input
        if axes is not None:
            axes_name = "reduce_axes"
            axes_np = np.asarray(axes, dtype=np.int64)
            axes_init = helper.make_tensor(
                name=axes_name,
                data_type=TensorProto.INT64,
                dims=axes_np.shape,
                vals=axes_np,
            )
            initializers.append(axes_init)
            inputs.append(axes_name)

        # Determine input and output shapes
        if not axes and not noop_with_empty_axes:
            outshape = np.sum(data, axis=None, keepdims=keepdims).shape
        elif not axes and noop_with_empty_axes:
            outshape = inshape
        else:
            outshape = np.sum(data, axis=axes, keepdims=keepdims).shape

        if dynamic:
            in_list = ["?"] * len(inshape)
            out_list = ["?"] * len(outshape)
        else:
            in_list = list(inshape)
            out_list = list(outshape)

        # Make a model node
        node = helper.make_node(
            func,
            inputs=inputs,
            outputs=["y"],
            keepdims=keepdims,
            noop_with_empty_axes=noop_with_empty_axes,
        )

        # Make a model graph and a model
        graph = helper.make_graph(
            [node],
            "reduce18_test",
            inputs=[helper.make_tensor_value_info("x", TensorProto.FLOAT, in_list)],
            initializer=initializers,
            outputs=[helper.make_tensor_value_info("y", TensorProto.FLOAT, out_list)],
        )
        model = helper.make_model(graph, producer_name="reduce18_test")

        inputs_dict = {"x": data}
        check_correctness(model, inputs_dict, opset=opset, rtol=1e-4, atol=1e-4)

    # Verify
    for keepdims in [True, False]:
        # no `axes` input && `noop_with_empty_axes` = 0 -> reduce over all dimensions.
        verify_reduce_func(
            func,
            np.random.randn(3, 2, 2).astype(np.float32),
            axes=[],
            keepdims=keepdims,
            noop_with_empty_axes=False,
        )

        # no `axes` input && `noop_with_empty_axes` = 0 -> reduce over all dimensions.
        verify_reduce_func(
            func,
            np.random.randn(3, 2, 2).astype(np.float32),
            axes=None,
            keepdims=keepdims,
            noop_with_empty_axes=False,
        )

        # no `axes` input && `noop_with_empty_axes` = 1 -> return the input unchanged.
        verify_reduce_func(
            func,
            np.random.randn(4, 3).astype(np.float32),
            axes=[],
            keepdims=keepdims,
            noop_with_empty_axes=True,
        )

        # no `axes` input && `noop_with_empty_axes` = 1 -> return the input unchanged.
        # (onnxruntime bug) Runtime error on the onnxruntime part
        # verify_reduce_func(
        #     func,
        #     np.random.randn(4, 3).astype(np.float32),
        #     axes=None,
        #     keepdims=keepdims,
        #     noop_with_empty_axes=True,
        # )

        # `axes` provided -> reduce over specified axes.
        verify_reduce_func(
            func,
            np.random.randn(3, 3, 3, 1).astype(np.float32),
            axes=(1, 2),
            keepdims=keepdims,
        )


def test_composite_reduce_funcs_axes_input_ir():
    for func in COMPOSITE_REDUCE_FUNCS:
        for keepdims in [True, False]:
            for dynamic in [True, False]:
                for input_shape, axes, noop_with_empty_axes in REDUCE_AXES_INPUT_TEST_CASES:
                    expected = _make_composite_reduce_expected_ir(
                        func,
                        input_shape,
                        axes,
                        noop_with_empty_axes,
                        keepdims,
                        dynamic,
                        axes_as_input=True,
                    )
                    verify_composite_reduce_axes_input_ir(
                        func,
                        input_shape,
                        axes,
                        noop_with_empty_axes,
                        keepdims,
                        dynamic,
                        18,
                        expected,
                    )


@pytest.mark.parametrize("in_dtype", [np.float32, np.int32])
@pytest.mark.parametrize("axis", [None, 0, 1, 2])
@pytest.mark.parametrize("keepdims", [None, True, False])
def test_arg_min_max(in_dtype, axis, keepdims):
    def verify_arg_min_max(input_dim, in_dtype, op_name="ArgMax", axis=None, keepdims=None):
        a_np1 = np.random.uniform(-10, 10, input_dim).astype(in_dtype)
        out_shape = list(a_np1.shape)
        def_axis = axis if axis is not None else 0
        if keepdims == 1 or keepdims is None:
            out_shape[def_axis] = 1
        else:
            out_shape.pop(def_axis)

        node = helper.make_node(op_name, inputs=["a_np1"], outputs=["out"])

        if keepdims is not None:
            keepdims_attr = helper.make_attribute("keepdims", keepdims)
            node.attribute.append(keepdims_attr)
        if axis is not None:
            axis_attr = helper.make_attribute("axis", axis)
            node.attribute.append(axis_attr)

        graph = helper.make_graph(
            [node],
            "argreduce_test",
            inputs=[helper.make_tensor_value_info("a_np1", TensorProto.INT32, list(a_np1.shape))],
            outputs=[helper.make_tensor_value_info("out", TensorProto.INT64, list(out_shape))],
        )

        model = helper.make_model(graph, producer_name="arg_min_max_test")
        check_correctness(model)

    verify_arg_min_max([3, 4, 4], in_dtype, "ArgMax", axis, keepdims)
    verify_arg_min_max([3, 4, 4], in_dtype, "ArgMin", axis, keepdims)


@pytest.mark.parametrize("axis", [-1, 0, 1])
@pytest.mark.parametrize("largest", [True, False])
def test_topk(axis: int, largest: int):
    in_shape = [32, 32, 32]
    k_value = 4
    out_shape = in_shape
    out_shape[axis] = k_value
    k = make_constant_node("k", TensorProto.INT64, [1], [k_value])
    node = onnx.helper.make_node(
        "TopK",
        inputs=["data", "k"],
        outputs=["values", "indices"],
        axis=axis,
        largest=largest,
    )
    graph = helper.make_graph(
        [k, node],
        "topk_test",
        inputs=[helper.make_tensor_value_info("data", TensorProto.FLOAT, in_shape)],
        outputs=[
            helper.make_tensor_value_info("values", TensorProto.FLOAT, out_shape),
            helper.make_tensor_value_info("indices", TensorProto.INT64, out_shape),
        ],
    )
    model = helper.make_model(graph, producer_name="topk_test")

    check_correctness(model, check_dtypes=True)


def test_expand():
    def _assert_expand_ir(name, input_shape, target_shape, output_shape, expected):
        shape_array = np.array(target_shape)
        shape_node = onnx.helper.make_node(
            "Constant",
            inputs=[],
            outputs=["shape"],
            value=onnx.helper.make_tensor(
                name="const_tensor",
                data_type=onnx.TensorProto.INT64,
                dims=shape_array.shape,
                vals=shape_array.flatten().astype("int64"),
            ),
        )
        expand_node = helper.make_node("Expand", ["in", "shape"], ["out"])

        graph = helper.make_graph(
            [shape_node, expand_node],
            "expand_teint64st",
            inputs=[helper.make_tensor_value_info("in", TensorProto.FLOAT, input_shape)],
            outputs=[helper.make_tensor_value_info("out", TensorProto.FLOAT, output_shape)],
        )

        model = helper.make_model(graph, producer_name=name)
        tvm_model = from_onnx(model, keep_params_in_input=True)
        tvm.ir.assert_structural_equal(tvm_model, expected)

    def _assert_expand_dynamic_shapeexpr_ir(name, input_shape, shape_input_shape, expected):
        shape_node = onnx.helper.make_node("Shape", inputs=["in_2"], outputs=["shape"])
        expand_node = helper.make_node("Expand", ["in", "shape"], ["out"])
        graph = helper.make_graph(
            [shape_node, expand_node],
            "expand_test",
            inputs=[
                helper.make_tensor_value_info("in", TensorProto.FLOAT, input_shape),
                helper.make_tensor_value_info("in_2", TensorProto.FLOAT, shape_input_shape),
            ],
            outputs=[helper.make_tensor_value_info("out", TensorProto.FLOAT, shape_input_shape)],
        )

        model = helper.make_model(graph, producer_name=name)
        tvm_model = from_onnx(model, keep_params_in_input=True)
        tvm.ir.assert_structural_equal(tvm_model, expected)

    @I.ir_module
    class ExpectedSameRank:
        @R.function
        def main(in_: R.Tensor((3, 1), dtype="float32")) -> R.Tensor((3, 4), dtype="float32"):
            R.func_attr({"num_input": 1})
            with R.dataflow():
                gv: R.Tensor((3, 4), dtype="float32") = R.broadcast_to(in_, R.shape([3, 4]))
                R.output(gv)
            return gv

    @I.ir_module
    class ExpectedHigherRank:
        @R.function
        def main(in_: R.Tensor((3, 1), dtype="float32")) -> R.Tensor((1, 3, 4), dtype="float32"):
            R.func_attr({"num_input": 1})
            with R.dataflow():
                gv: R.Tensor((1, 3, 4), dtype="float32") = R.broadcast_to(in_, R.shape([1, 3, 4]))
                R.output(gv)
            return gv

    @I.ir_module
    class ExpectedSameSuffix:
        @R.function
        def main(in_: R.Tensor((3, 1), dtype="float32")) -> R.Tensor((1, 1, 3, 1), dtype="float32"):
            R.func_attr({"num_input": 1})
            with R.dataflow():
                gv: R.Tensor((1, 1, 3, 1), dtype="float32") = R.broadcast_to(
                    in_, R.shape([1, 1, 3, 1])
                )
                R.output(gv)
            return gv

    @I.ir_module
    class ExpectedDynamicShape:
        @R.function
        def main(
            in_: R.Tensor((1, 32, 32), dtype="float32"),
            in_2: R.Tensor(("batch", 32, 32), dtype="float32"),
        ) -> R.Tensor(("batch", 32, 32), dtype="float32"):
            batch = T.int64()
            R.func_attr({"num_input": 2})
            with R.dataflow():
                gv: R.Tensor((batch, 32, 32), dtype="float32") = R.broadcast_to(
                    in_, R.shape([batch, 32, 32])
                )
                R.output(gv)
            return gv

    _assert_expand_ir("expand_with_dim_unchanged_test", [3, 1], [3, 4], [3, 4], ExpectedSameRank)
    _assert_expand_ir("expand_with_diff_dim", [3, 1], [1, 3, 4], [1, 3, 4], ExpectedHigherRank)
    _assert_expand_ir(
        "expand_with_the_same_suffix_dims", [3, 1], [1, 1, 3, 1], [1, 1, 3, 1], ExpectedSameSuffix
    )
    _assert_expand_dynamic_shapeexpr_ir(
        "expand_with_dynamic_dim", [1, 32, 32], ["batch", 32, 32], ExpectedDynamicShape
    )


def test_expand_incompatible_broadcasting():
    """
    This test case reproduces the error where input tensor shape at dim 1 is 25
    and target shape at dim 3 is 56, which violates ONNX broadcasting rules
    """

    def _test_expand_error_case(name, data_shape, target_shape_vals):
        data = np.random.uniform(size=data_shape).astype(np.float32)

        shape_array = np.array(target_shape_vals, dtype=np.int64)
        shape_node = onnx.helper.make_node(
            "Constant",
            inputs=[],
            outputs=["shape"],
            value=onnx.helper.make_tensor(
                name="const_tensor",
                data_type=onnx.TensorProto.INT64,
                dims=shape_array.shape,
                vals=shape_array.flatten(),
            ),
        )

        expand_node = helper.make_node("Expand", ["in", "shape"], ["out"])

        graph = helper.make_graph(
            [shape_node, expand_node],
            "expand_error_test",
            inputs=[helper.make_tensor_value_info("in", TensorProto.FLOAT, list(data.shape))],
            outputs=[helper.make_tensor_value_info("out", TensorProto.FLOAT, target_shape_vals)],
        )

        model = helper.make_model(graph, producer_name=name)

        with pytest.raises(ValueError) as exc_info:
            from_onnx(model, keep_params_in_input=True)

        error_msg = str(exc_info.value)
        assert "broadcast" in error_msg.lower() or "incompatible" in error_msg.lower(), (
            f"Expected broadcasting error, but got: {error_msg}"
        )

    # Test case 1: Reproduce the exact error from the issue-17769
    # Input shape: (25,), target shape: (1, 1, 1, 56)
    # This should faill because input dim 1 (25) != target dim 3 (56) and neither is 1
    _test_expand_error_case(
        "expand_incompatible_25_to_56",
        data_shape=(25,),
        target_shape_vals=(1, 1, 1, 56),
    )

    # Test case 2: Another incompatible case
    # Input shape: (1, 25), target shape: (1, 1, 1, 56)
    # After right-alignment, input (1, 1, 1, 25) vs. target (1, 1, 1, 56)
    # This should fail because 25 != 56 and neither is 1
    _test_expand_error_case(
        "expand_incompatible_aligned_25_to_56",
        data_shape=(1, 25),
        target_shape_vals=(1, 1, 1, 56),
    )

    # Test case 3: Valid case for comparison - should not raise error
    def _test_expand_valid_case():
        """Test a valid expand case to ensure our fix doesn't break valid operations"""
        data_shape = (1, 25)
        target_shape_vals = [2, 25]  # Valid: input (1, 25) can broadcast to (2, 25)

        data = np.random.uniform(size=data_shape).astype(np.float32)
        shape_array = np.array(target_shape_vals, dtype=np.int64)

        shape_node = onnx.helper.make_node(
            "Constant",
            inputs=[],
            outputs=["shape"],
            value=onnx.helper.make_tensor(
                name="const_tensor",
                data_type=onnx.TensorProto.INT64,
                dims=shape_array.shape,
                vals=shape_array.flatten(),
            ),
        )

        expand_node = helper.make_node("Expand", ["in", "shape"], ["out"])

        graph = helper.make_graph(
            [shape_node, expand_node],
            "expand_valid_test",
            inputs=[helper.make_tensor_value_info("in", TensorProto.FLOAT, list(data.shape))],
            outputs=[helper.make_tensor_value_info("out", TensorProto.FLOAT, target_shape_vals)],
        )

        model = helper.make_model(graph, producer_name="expand_valid_test_case")

        try:
            tvm_model = from_onnx(model, keep_params_in_input=True)
        except Exception as e:
            pytest.fail(f"Valid expand case should not fail, but got error: {e}")

    _test_expand_valid_case()


# TODO(jwfromm) Current approach to dynamic expand is technically not well formed. Reenable once fixed.
@pytest.mark.skip("Produces ill-formed IR")
def test_constantofshape():
    def verify_constantofshape(input_dim, value, dtype):
        fill_node = helper.make_node(
            "ConstantOfShape",
            ["input"],
            ["output"],
            value=helper.make_tensor(
                "value", helper.np_dtype_to_tensor_dtype(np.dtype(dtype)), (1,), (value,)
            ),
        )

        inputs = [helper.make_tensor_value_info("input", TensorProto.INT64, [len(input_dim)])]

        graph = helper.make_graph(
            [fill_node],
            "fill_test",
            inputs,
            initializer=[
                helper.make_tensor(
                    "input",
                    TensorProto.INT64,
                    [len(input_dim)],
                    np.asarray(input_dim).astype("int64"),
                )
            ],
            outputs=[
                helper.make_tensor_value_info(
                    "output", helper.np_dtype_to_tensor_dtype(np.dtype(dtype)), input_dim
                )
            ],
        )

        model = helper.make_model(graph, producer_name="fill_test")
        tvm_model = from_onnx(model, keep_params_in_input=True)
        assert tuple(dim.value for dim in tvm_model["main"].ret_ty.shape.values) == input_dim

    verify_constantofshape((2, 3, 4, 5), 10, "float32")
    verify_constantofshape((3, 3), 0, "int32")
    verify_constantofshape((1, 2, 3), -1, "float32")


def test_constantofshape_default_value():
    """ConstantOfShape value attribute should default to float32 zero."""
    shape_init = helper.make_tensor("shape", TensorProto.INT64, [2], [2, 3])
    node = helper.make_node("ConstantOfShape", ["shape"], ["y"])
    graph = helper.make_graph(
        [node],
        "constantofshape_default_value_test",
        inputs=[],
        outputs=[helper.make_tensor_value_info("y", TensorProto.FLOAT, None)],
        initializer=[shape_init],
    )
    model = helper.make_model(graph, producer_name="constantofshape_default_value_test")

    tvm_model = from_onnx(model)

    @I.ir_module
    class Expected:
        @R.function
        def main() -> R.Tensor((2, 3), dtype="float32"):
            R.func_attr({"num_input": 0})
            with R.dataflow():
                gv: R.Tensor((2, 3), dtype="float32") = R.broadcast_to(
                    R.const(0.0, "float32"), R.shape([2, 3])
                )
                R.output(gv)
            return gv

    tvm.ir.assert_structural_equal(tvm_model, Expected)


def test_slice():
    def verify_slice(data_shape, output_shape, starts, ends, expected, axes=None, steps=None):
        if isinstance(starts, list):
            starts = np.array(starts, "int64")
        if isinstance(ends, list):
            ends = np.array(ends, "int64")
        if isinstance(axes, list):
            axes = np.array(axes, "int64")
        if isinstance(steps, list):
            steps = np.array(steps, "int64")

        slice_inputs = ["x", "starts", "ends"]
        initializer = [
            helper.make_tensor("starts", TensorProto.INT64, starts.shape, starts),
            helper.make_tensor("ends", TensorProto.INT64, ends.shape, ends),
        ]

        if axes is not None:
            initializer.append(helper.make_tensor("axes", TensorProto.INT64, axes.shape, axes))
            slice_inputs.append("axes")
        if steps is not None:
            initializer.append(helper.make_tensor("steps", TensorProto.INT64, steps.shape, steps))
            slice_inputs.append("steps")

        slice_node = helper.make_node("Slice", inputs=slice_inputs, outputs=["y"])

        graph = helper.make_graph(
            [slice_node],
            "slice_test",
            inputs=[
                helper.make_tensor_value_info("x", TensorProto.FLOAT, data_shape),
            ],
            outputs=[helper.make_tensor_value_info("y", TensorProto.FLOAT, output_shape)],
            initializer=initializer,
        )

        model = helper.make_model(graph, producer_name="slice_test")
        tvm_model = from_onnx(model, keep_params_in_input=True)
        tvm_model["main"] = tvm_model["main"].without_attr("params")
        tvm.ir.assert_structural_equal(tvm_model, expected)

    @I.ir_module
    class ExpectedSliceAxesAndSteps:
        @R.function
        def main(
            x: R.Tensor((20, 10, 5), dtype="float32"),
            starts: R.Tensor((2,), dtype="int64"),
            ends: R.Tensor((2,), dtype="int64"),
            axes: R.Tensor((2,), dtype="int64"),
            steps: R.Tensor((2,), dtype="int64"),
        ) -> R.Tensor((3, 10, 5), dtype="float32"):
            R.func_attr({"num_input": 1})
            with R.dataflow():
                gv: R.Tensor((3, 10, 5), dtype="float32") = R.strided_slice(
                    x,
                    axes=[0, 1],
                    begin=[0, 0],
                    end=[3, 10],
                    strides=[1, 1],
                    assume_inbound=False,
                )
                R.output(gv)
            return gv

    @I.ir_module
    class ExpectedSliceDefaultAxesAndSteps:
        @R.function
        def main(
            x: R.Tensor((20, 10, 5), dtype="float32"),
            starts: R.Tensor((2,), dtype="int64"),
            ends: R.Tensor((2,), dtype="int64"),
        ) -> R.Tensor((3, 10, 5), dtype="float32"):
            R.func_attr({"num_input": 1})
            with R.dataflow():
                gv: R.Tensor((3, 10, 5), dtype="float32") = R.strided_slice(
                    x,
                    axes=[0, 1],
                    begin=[0, 0],
                    end=[3, 10],
                    strides=[1, 1],
                    assume_inbound=False,
                )
                R.output(gv)
            return gv

    @I.ir_module
    class ExpectedSliceNegativeSteps:
        @R.function
        def main(
            x: R.Tensor((20, 10, 5), dtype="float32"),
            starts: R.Tensor((3,), dtype="int64"),
            ends: R.Tensor((3,), dtype="int64"),
            axes: R.Tensor((3,), dtype="int64"),
            steps: R.Tensor((3,), dtype="int64"),
        ) -> R.Tensor((19, 3, 2), dtype="float32"):
            R.func_attr({"num_input": 1})
            with R.dataflow():
                gv: R.Tensor((19, 3, 2), dtype="float32") = R.strided_slice(
                    x,
                    axes=[0, 1, 2],
                    begin=[20, 10, 4],
                    end=[0, 0, 1],
                    strides=[-1, -3, -2],
                    assume_inbound=False,
                )
                R.output(gv)
            return gv

    @I.ir_module
    class ExpectedSliceAxesOnly:
        @R.function
        def main(
            x: R.Tensor((20, 10, 5), dtype="float32"),
            starts: R.Tensor((2,), dtype="int64"),
            ends: R.Tensor((2,), dtype="int64"),
            axes: R.Tensor((2,), dtype="int64"),
        ) -> R.Tensor((20, 3, 5), dtype="float32"):
            R.func_attr({"num_input": 1})
            with R.dataflow():
                gv: R.Tensor((20, 3, 5), dtype="float32") = R.strided_slice(
                    x,
                    axes=[1, 2],
                    begin=[0, 0],
                    end=[3, 10],
                    strides=[1, 1],
                    assume_inbound=False,
                )
                R.output(gv)
            return gv

    # Test with all parameters set.
    verify_slice(
        [20, 10, 5],
        [3, 10, 5],
        starts=[0, 0],
        ends=[3, 10],
        axes=[0, 1],
        steps=[1, 1],
        expected=ExpectedSliceAxesAndSteps,
    )
    # Test with default axes and steps.
    verify_slice(
        [20, 10, 5],
        [3, 10, 5],
        starts=[0, 0],
        ends=[3, 10],
        expected=ExpectedSliceDefaultAxesAndSteps,
    )
    # Test with negative steps.
    verify_slice(
        [20, 10, 5],
        [19, 3, 2],
        starts=[20, 10, 4],  # NOTE: the start is out of bounds
        ends=[0, 0, 1],
        steps=[-1, -3, -2],
        axes=[0, 1, 2],
        expected=ExpectedSliceNegativeSteps,
    )
    verify_slice(
        [20, 10, 5],
        [10, 5],
        starts=[0, 0],
        ends=[3, 10],
        axes=[1, 2],
        expected=ExpectedSliceAxesOnly,
    )
    verify_slice(
        [20, 10, 5],
        [10, 5],
        starts=[0, 0],
        ends=[3, 10],
        axes=[1, 2],
        expected=ExpectedSliceAxesOnly,
    )

    # TODO (gigiblender): Enable this test when we have a way to pass the steps but not axes.
    # verify_slice(
    #     [20, 10, 5],
    #     [19, 3, 2],
    #     starts=[20, 10, 4],
    #     ends=[0, 0, 1],
    #     steps=[-1, -3, -2],
    # )


def test_slice_dynamic_inputs_ir():
    slice_node = helper.make_node("Slice", ["x", "starts", "ends", "axes", "steps"], ["y"])

    graph = helper.make_graph(
        [slice_node],
        "slice_dynamic_inputs_ir",
        inputs=[
            helper.make_tensor_value_info("x", TensorProto.FLOAT, [20, 10, 5]),
            helper.make_tensor_value_info("starts", TensorProto.INT64, [2]),
            helper.make_tensor_value_info("ends", TensorProto.INT64, [2]),
            helper.make_tensor_value_info("axes", TensorProto.INT64, [2]),
            helper.make_tensor_value_info("steps", TensorProto.INT64, [2]),
        ],
        outputs=[helper.make_tensor_value_info("y", TensorProto.FLOAT, [3, 10, 5])],
    )

    model = helper.make_model(graph, producer_name="slice_dynamic_inputs_ir_test")
    tvm_model = from_onnx(model, opset=13, keep_params_in_input=True)

    @I.ir_module
    class Expected:
        @R.function
        def main(
            x: R.Tensor((20, 10, 5), dtype="float32"),
            starts: R.Tensor((2,), dtype="int64"),
            ends: R.Tensor((2,), dtype="int64"),
            axes: R.Tensor((2,), dtype="int64"),
            steps: R.Tensor((2,), dtype="int64"),
        ) -> R.Tensor(dtype="float32", ndim=3):
            R.func_attr({"num_input": 5})
            with R.dataflow():
                lv: R.Tensor((2,), dtype="bool") = R.less(axes, R.const(0, "int64"))
                lv1: R.Tensor((2,), dtype="int64") = R.add(axes, R.const(3, "int64"))
                lv2: R.Shape([20, 10, 5]) = R.shape_of(x)
                lv3: R.Tensor((2,), dtype="int64") = R.where(lv, lv1, axes)
                lv4: R.Tensor((3,), dtype="int64") = R.shape_to_tensor(lv2)
                lv5: R.Tensor((3,), dtype="int64") = R.scatter_elements(
                    R.const([0, 0, 0], "int64"), lv3, starts, axis=0, reduction="update"
                )
                lv6: R.Tensor((3,), dtype="int64") = R.scatter_elements(
                    lv4, lv3, ends, axis=0, reduction="update"
                )
                lv7: R.Tensor((3,), dtype="int64") = R.scatter_elements(
                    R.const([1, 1, 1], "int64"), lv3, steps, axis=0, reduction="update"
                )
                gv: R.Tensor(dtype="float32", ndim=3) = R.dynamic_strided_slice(x, lv5, lv6, lv7)
                R.output(gv)
            return gv

    tvm.ir.assert_structural_equal(tvm_model, Expected)


def test_slice_dynamic_inputs_length_validation():
    slice_node = helper.make_node("Slice", ["x", "starts", "ends", "axes", "steps"], ["y"])

    graph = helper.make_graph(
        [slice_node],
        "slice_dynamic_inputs_length_validation",
        inputs=[
            helper.make_tensor_value_info("x", TensorProto.FLOAT, [20, 10, 5]),
            helper.make_tensor_value_info("starts", TensorProto.INT64, [2]),
            helper.make_tensor_value_info("ends", TensorProto.INT64, [1]),
            helper.make_tensor_value_info("axes", TensorProto.INT64, [2]),
            helper.make_tensor_value_info("steps", TensorProto.INT64, [2]),
        ],
        outputs=[helper.make_tensor_value_info("y", TensorProto.FLOAT, [3, 10, 5])],
    )

    model = helper.make_model(graph, producer_name="slice_dynamic_inputs_length_validation_test")
    with pytest.raises(ValueError, match="starts and ends to have the same length"):
        from_onnx(model, opset=13, keep_params_in_input=True)


def test_slice_dynamic_shape_expr_input_validation():
    shape_node = helper.make_node("Shape", ["x"], ["y"])
    slice_node = helper.make_node("Slice", ["y", "starts", "ends", "axes", "steps"], ["z"])

    graph = helper.make_graph(
        [shape_node, slice_node],
        "slice_dynamic_shape_expr_input_validation",
        inputs=[
            helper.make_tensor_value_info("x", TensorProto.FLOAT, [20, 10, 5]),
            helper.make_tensor_value_info("starts", TensorProto.INT64, [1]),
            helper.make_tensor_value_info("ends", TensorProto.INT64, [1]),
            helper.make_tensor_value_info("axes", TensorProto.INT64, [1]),
            helper.make_tensor_value_info("steps", TensorProto.INT64, [1]),
        ],
        outputs=[helper.make_tensor_value_info("z", TensorProto.INT64, [1])],
    )

    model = helper.make_model(graph, producer_name="slice_dynamic_shape_expr_input_validation_test")
    with pytest.raises(ValueError, match="does not support ShapeExpr input"):
        from_onnx(model, opset=13, keep_params_in_input=True)


def test_slice_zero_step_validation():
    slice_node = helper.make_node("Slice", ["x", "starts", "ends", "axes", "steps"], ["y"])

    graph = helper.make_graph(
        [slice_node],
        "slice_zero_step_validation",
        inputs=[helper.make_tensor_value_info("x", TensorProto.FLOAT, [20, 10, 5])],
        initializer=[
            helper.make_tensor("starts", TensorProto.INT64, [2], vals=[0, 0]),
            helper.make_tensor("ends", TensorProto.INT64, [2], vals=[3, 10]),
            helper.make_tensor("axes", TensorProto.INT64, [2], vals=[0, 1]),
            helper.make_tensor("steps", TensorProto.INT64, [2], vals=[1, 0]),
        ],
        outputs=[helper.make_tensor_value_info("y", TensorProto.FLOAT, [3, 10, 5])],
    )

    model = helper.make_model(graph, producer_name="slice_zero_step_validation_test")
    with pytest.raises(ValueError, match="step values must be non-zero"):
        from_onnx(model, opset=13)


def test_slice_dynamic_shape():
    def verify_slice(data_shape, output_shape, starts, ends, axes, expected):
        if isinstance(starts, list):
            starts = np.array(starts, "int64")
        if isinstance(ends, list):
            ends = np.array(ends, "int64")
        if isinstance(axes, list):
            axes = np.array(axes, "int64")

        slice_inputs = ["y", "starts", "ends"]
        initializer = [
            helper.make_tensor("starts", TensorProto.INT64, starts.shape, starts),
            helper.make_tensor("ends", TensorProto.INT64, ends.shape, ends),
            helper.make_tensor("axes", TensorProto.INT64, axes.shape, axes),
        ]
        slice_inputs.append("axes")

        shape_node = helper.make_node("Shape", inputs=["x"], outputs=["y"])
        slice_node = helper.make_node("Slice", inputs=slice_inputs, outputs=["z"])

        graph = helper.make_graph(
            [shape_node, slice_node],
            "slice_test",
            inputs=[
                helper.make_tensor_value_info("x", TensorProto.FLOAT, data_shape),
            ],
            outputs=[helper.make_tensor_value_info("z", TensorProto.INT64, output_shape)],
            initializer=initializer,
        )

        model = helper.make_model(graph, producer_name="slice_test")
        tvm_model = from_onnx(model, keep_params_in_input=True)
        assert len(tvm_model["main"].attrs["params"]) == 3
        tvm_model["main"] = tvm_model["main"].without_attr("params")
        tvm.ir.assert_structural_equal(tvm_model, expected)

    @I.ir_module
    class ExpectedShapeSlice0:
        @R.function
        def main(
            x: R.Tensor((20, 10, 5), dtype="float32"),
            starts: R.Tensor((1,), dtype="int64"),
            ends: R.Tensor((1,), dtype="int64"),
            axes: R.Tensor((1,), dtype="int64"),
        ) -> R.Tensor((2,), dtype="int64"):
            R.func_attr({"num_input": 1})
            with R.dataflow():
                gv: R.Tensor((2,), dtype="int64") = R.const([20, 10], "int64")
                R.output(gv)
            return gv

    @I.ir_module
    class ExpectedShapeSlice1:
        @R.function
        def main(
            x: R.Tensor(("A", 10, 5), dtype="float32"),
            starts: R.Tensor((1,), dtype="int64"),
            ends: R.Tensor((1,), dtype="int64"),
            axes: R.Tensor((1,), dtype="int64"),
        ) -> R.Shape(ndim=2):
            A = T.int64()
            R.func_attr({"num_input": 1})
            with R.dataflow():
                gv: R.Shape([A, 10]) = R.shape([A, 10])
                R.output(gv)
            return gv

    @I.ir_module
    class ExpectedShapeSlice2:
        @R.function
        def main(
            x: R.Tensor(("A", "B", 5), dtype="float32"),
            starts: R.Tensor((1,), dtype="int64"),
            ends: R.Tensor((1,), dtype="int64"),
            axes: R.Tensor((1,), dtype="int64"),
        ) -> R.Shape(ndim=2):
            A = T.int64()
            B = T.int64()
            R.func_attr({"num_input": 1})
            with R.dataflow():
                gv: R.Shape([A, B]) = R.shape([A, B])
                R.output(gv)
            return gv

    @I.ir_module
    class ExpectedShapeSlice3:
        @R.function
        def main(
            x: R.Tensor((20, 10, "C"), dtype="float32"),
            starts: R.Tensor((1,), dtype="int64"),
            ends: R.Tensor((1,), dtype="int64"),
            axes: R.Tensor((1,), dtype="int64"),
        ) -> R.Tensor((2,), dtype="int64"):
            C = T.int64()
            R.func_attr({"num_input": 1})
            with R.dataflow():
                gv: R.Tensor((2,), dtype="int64") = R.const([20, 10], "int64")
                R.output(gv)
            return gv

    @I.ir_module
    class ExpectedShapeSlice4:
        @R.function
        def main(
            x: R.Tensor(("A", "B", "C"), dtype="float32"),
            starts: R.Tensor((1,), dtype="int64"),
            ends: R.Tensor((1,), dtype="int64"),
            axes: R.Tensor((1,), dtype="int64"),
        ) -> R.Shape(ndim=2):
            A = T.int64()
            B = T.int64()
            C = T.int64()
            R.func_attr({"num_input": 1})
            with R.dataflow():
                gv: R.Shape([A, B]) = R.shape([A, B])
                R.output(gv)
            return gv

    @I.ir_module
    class ExpectedShapeSlice5:
        @R.function
        def main(
            x: R.Tensor((20, 10, 5), dtype="float32"),
            starts: R.Tensor((1,), dtype="int64"),
            ends: R.Tensor((1,), dtype="int64"),
            axes: R.Tensor((1,), dtype="int64"),
        ) -> R.Tensor((1,), dtype="int64"):
            R.func_attr({"num_input": 1})
            with R.dataflow():
                gv: R.Tensor((1,), dtype="int64") = R.const([10], "int64")
                R.output(gv)
            return gv

    @I.ir_module
    class ExpectedShapeSlice6:
        @R.function
        def main(
            x: R.Tensor(("A", 10, 5), dtype="float32"),
            starts: R.Tensor((1,), dtype="int64"),
            ends: R.Tensor((1,), dtype="int64"),
            axes: R.Tensor((1,), dtype="int64"),
        ) -> R.Tensor((1,), dtype="int64"):
            A = T.int64()
            R.func_attr({"num_input": 1})
            with R.dataflow():
                gv: R.Tensor((1,), dtype="int64") = R.const([10], "int64")
                R.output(gv)
            return gv

    @I.ir_module
    class ExpectedShapeSlice7:
        @R.function
        def main(
            x: R.Tensor(("A", "B", 5), dtype="float32"),
            starts: R.Tensor((1,), dtype="int64"),
            ends: R.Tensor((1,), dtype="int64"),
            axes: R.Tensor((1,), dtype="int64"),
        ) -> R.Shape(ndim=1):
            A = T.int64()
            B = T.int64()
            R.func_attr({"num_input": 1})
            with R.dataflow():
                gv: R.Shape([B]) = R.shape([B])
                R.output(gv)
            return gv

    @I.ir_module
    class ExpectedShapeSlice8:
        @R.function
        def main(
            x: R.Tensor((20, 10, "C"), dtype="float32"),
            starts: R.Tensor((1,), dtype="int64"),
            ends: R.Tensor((1,), dtype="int64"),
            axes: R.Tensor((1,), dtype="int64"),
        ) -> R.Tensor((1,), dtype="int64"):
            C = T.int64()
            R.func_attr({"num_input": 1})
            with R.dataflow():
                gv: R.Tensor((1,), dtype="int64") = R.const([10], "int64")
                R.output(gv)
            return gv

    @I.ir_module
    class ExpectedShapeSlice9:
        @R.function
        def main(
            x: R.Tensor(("A", "B", "C"), dtype="float32"),
            starts: R.Tensor((1,), dtype="int64"),
            ends: R.Tensor((1,), dtype="int64"),
            axes: R.Tensor((1,), dtype="int64"),
        ) -> R.Shape(ndim=1):
            A = T.int64()
            B = T.int64()
            C = T.int64()
            R.func_attr({"num_input": 1})
            with R.dataflow():
                gv: R.Shape([B]) = R.shape([B])
                R.output(gv)
            return gv

    @I.ir_module
    class ExpectedShapeSlice10:
        @R.function
        def main(
            x: R.Tensor((20, 10, 5), dtype="float32"),
            starts: R.Tensor((1,), dtype="int64"),
            ends: R.Tensor((1,), dtype="int64"),
            axes: R.Tensor((1,), dtype="int64"),
        ) -> R.Tensor((2,), dtype="int64"):
            R.func_attr({"num_input": 1})
            with R.dataflow():
                gv: R.Tensor((2,), dtype="int64") = R.const([10, 5], "int64")
                R.output(gv)
            return gv

    @I.ir_module
    class ExpectedShapeSlice11:
        @R.function
        def main(
            x: R.Tensor(("A", 10, 5), dtype="float32"),
            starts: R.Tensor((1,), dtype="int64"),
            ends: R.Tensor((1,), dtype="int64"),
            axes: R.Tensor((1,), dtype="int64"),
        ) -> R.Tensor((2,), dtype="int64"):
            A = T.int64()
            R.func_attr({"num_input": 1})
            with R.dataflow():
                gv: R.Tensor((2,), dtype="int64") = R.const([10, 5], "int64")
                R.output(gv)
            return gv

    @I.ir_module
    class ExpectedShapeSlice12:
        @R.function
        def main(
            x: R.Tensor(("A", "B", 5), dtype="float32"),
            starts: R.Tensor((1,), dtype="int64"),
            ends: R.Tensor((1,), dtype="int64"),
            axes: R.Tensor((1,), dtype="int64"),
        ) -> R.Shape(ndim=2):
            A = T.int64()
            B = T.int64()
            R.func_attr({"num_input": 1})
            with R.dataflow():
                gv: R.Shape([B, 5]) = R.shape([B, 5])
                R.output(gv)
            return gv

    @I.ir_module
    class ExpectedShapeSlice13:
        @R.function
        def main(
            x: R.Tensor((20, 10, "C"), dtype="float32"),
            starts: R.Tensor((1,), dtype="int64"),
            ends: R.Tensor((1,), dtype="int64"),
            axes: R.Tensor((1,), dtype="int64"),
        ) -> R.Shape(ndim=2):
            C = T.int64()
            R.func_attr({"num_input": 1})
            with R.dataflow():
                gv: R.Shape([10, C]) = R.shape([10, C])
                R.output(gv)
            return gv

    @I.ir_module
    class ExpectedShapeSlice14:
        @R.function
        def main(
            x: R.Tensor(("A", "B", "C"), dtype="float32"),
            starts: R.Tensor((1,), dtype="int64"),
            ends: R.Tensor((1,), dtype="int64"),
            axes: R.Tensor((1,), dtype="int64"),
        ) -> R.Shape(ndim=2):
            A = T.int64()
            B = T.int64()
            C = T.int64()
            R.func_attr({"num_input": 1})
            with R.dataflow():
                gv: R.Shape([B, C]) = R.shape([B, C])
                R.output(gv)
            return gv

    verify_slice([20, 10, 5], [2], starts=[0], ends=[2], axes=[0], expected=ExpectedShapeSlice0)
    verify_slice(["A", 10, 5], [2], starts=[0], ends=[2], axes=[0], expected=ExpectedShapeSlice1)
    verify_slice(["A", "B", 5], [2], starts=[0], ends=[2], axes=[0], expected=ExpectedShapeSlice2)
    verify_slice([20, 10, "C"], [2], starts=[0], ends=[2], axes=[0], expected=ExpectedShapeSlice3)
    verify_slice(["A", "B", "C"], [2], starts=[0], ends=[2], axes=[0], expected=ExpectedShapeSlice4)
    verify_slice([20, 10, 5], [1], starts=[1], ends=[2], axes=[0], expected=ExpectedShapeSlice5)
    verify_slice(["A", 10, 5], [1], starts=[1], ends=[2], axes=[0], expected=ExpectedShapeSlice6)
    verify_slice(["A", "B", 5], [1], starts=[1], ends=[2], axes=[0], expected=ExpectedShapeSlice7)
    verify_slice([20, 10, "C"], [1], starts=[1], ends=[2], axes=[0], expected=ExpectedShapeSlice8)
    verify_slice(["A", "B", "C"], [1], starts=[1], ends=[2], axes=[0], expected=ExpectedShapeSlice9)
    verify_slice([20, 10, 5], [2], starts=[1], ends=[3], axes=[0], expected=ExpectedShapeSlice10)
    verify_slice(["A", 10, 5], [2], starts=[1], ends=[3], axes=[0], expected=ExpectedShapeSlice11)
    verify_slice(["A", "B", 5], [2], starts=[1], ends=[3], axes=[0], expected=ExpectedShapeSlice12)
    verify_slice([20, 10, "C"], [2], starts=[1], ends=[3], axes=[0], expected=ExpectedShapeSlice13)
    verify_slice(
        ["A", "B", "C"], [2], starts=[1], ends=[3], axes=[0], expected=ExpectedShapeSlice14
    )


# TODO Enable dynamism
@pytest.mark.parametrize("dynamic", [False])
def test_attention(dynamic):
    def verify_attention(
        input_,
        weight,
        bias,
        mask_index,
        num_heads,
        mask_filter_value,
        qkv_hidden_sizes,
        relative_position_bias,
    ):
        node = onnx.helper.make_node(
            "Attention",
            inputs=["input", "weight", "bias", "mask_index", "", "relative_position_bias"],
            outputs=["output"],
            domain="com.microsoft",
            num_heads=num_heads,
            # TODO(jwfromm) OnnxRT doesnt work with this attribute, figure out why not.
            # mask_filter_value=mask_filter_value,
            qkv_hidden_sizes=qkv_hidden_sizes,
        )

        input_shape = list(input_.shape)
        weight_shape = list(weight.shape)
        bias_shape = list(bias.shape)
        mask_shape = list(mask_index.shape)
        relative_position_bias_shape = list(relative_position_bias.shape)
        output_shape = list(input_.shape)
        if dynamic:
            input_shape = ["?" for _ in range(len(input_.shape))]
            weight_shape = ["?" for _ in range(len(weight.shape))]
            bias_shape = ["?" for _ in range(len(bias.shape))]
            mask_shape = ["?" for _ in range(len(mask_index.shape))]
            output_shape = ["?" for _ in range(len(input_.shape))]

        graph = helper.make_graph(
            [node],
            "attention_test",
            inputs=[
                helper.make_tensor_value_info("input", TensorProto.FLOAT, input_shape),
                helper.make_tensor_value_info("weight", TensorProto.FLOAT, weight_shape),
                helper.make_tensor_value_info("bias", TensorProto.FLOAT, bias_shape),
                helper.make_tensor_value_info("mask_index", TensorProto.INT32, mask_shape),
                helper.make_tensor_value_info(
                    "relative_position_bias", TensorProto.FLOAT, relative_position_bias_shape
                ),
            ],
            outputs=[
                helper.make_tensor_value_info("output", TensorProto.FLOAT, output_shape),
            ],
        )

        model = helper.make_model(graph, producer_name="attention_test")
        tvm_model = from_onnx(model, keep_params_in_input=True)

        @I.ir_module
        class ExpectedAttention:
            @R.function
            def main(
                input: R.Tensor((4, 4, 128), dtype="float32"),
                weight: R.Tensor((128, 480), dtype="float32"),
                bias: R.Tensor((480,), dtype="float32"),
                mask_index: R.Tensor((4, 4), dtype="int32"),
                relative_position_bias: R.Tensor((4, 12, 4, 4), dtype="float32"),
            ) -> R.Tensor((4, 4, 96), dtype="float32"):
                R.func_attr({"num_input": 5})
                with R.dataflow():
                    lv: R.Tensor((4, 4), dtype="int32") = R.subtract(
                        R.const(1, "int32"), mask_index
                    )
                    lv1: R.Tensor((4, 4), dtype="float32") = R.astype(lv, dtype="float32")
                    lv2: R.Tensor((4, 4), dtype="float32") = R.multiply(
                        lv1, R.const(-10000.0, "float32")
                    )
                    lv3: R.Tensor((4, 1, 1, 4), dtype="float32") = R.reshape(
                        lv2, R.shape([4, 1, 1, 4])
                    )
                    lv4: R.Tensor((4, 4, 480), dtype="float32") = R.matmul(input, weight)
                    lv5: R.Tensor((4, 4, 480), dtype="float32") = R.add(lv4, bias)
                    lv6: R.Tuple(
                        R.Tensor((4, 4, 192), dtype="float32"),
                        R.Tensor((4, 4, 192), dtype="float32"),
                        R.Tensor((4, 4, 96), dtype="float32"),
                    ) = R.split(lv5, indices_or_sections=[192, 384], axis=2)
                    lv7: R.Tensor((4, 4, 192), dtype="float32") = lv6[0]
                    lv8: R.Tensor((4, 4, 192), dtype="float32") = lv6[1]
                    lv9: R.Tensor((4, 4, 96), dtype="float32") = lv6[2]
                    lv10: R.Tensor((4, 4, 12, 16), dtype="float32") = R.reshape(
                        lv7, R.shape([4, 4, 12, 16])
                    )
                    lv11: R.Tensor((4, 4, 12, 16), dtype="float32") = R.reshape(
                        lv8, R.shape([4, 4, 12, 16])
                    )
                    lv12: R.Tensor((4, 4, 12, 8), dtype="float32") = R.reshape(
                        lv9, R.shape([4, 4, 12, 8])
                    )
                    lv13: R.Tensor((4, 12, 4, 4), dtype="float32") = R.add(
                        relative_position_bias, lv3
                    )
                    lv14: R.Tensor((4, 4, 12, 8), dtype="float32") = R.nn.attention(
                        lv10, lv11, lv12, lv13
                    )
                    lv15: R.Tensor((4, 4, 96), dtype="float32") = R.reshape(
                        lv14, R.shape([4, 4, 96])
                    )
                    gv: R.Tensor((4, 4, 96), dtype="float32") = lv15
                    R.output(gv)
                return gv

        tvm.ir.assert_structural_equal(tvm_model, ExpectedAttention)
        # "present" output should be nullptr when the "past" input isn't included,
        # but ort requires an output shape to be specified?
        # verify_with_ort_with_inputs(
        #     model,
        #     [input_, weight, bias, mask_index],
        #     [input_.shape, present_output_shape],
        #     target=target,
        #     dev=dev,
        #     rtol=1e-4,
        #     atol=1e-4,
        # )

    input_hidden_size = 128
    batch_size = 4
    sequence_length = 4
    num_heads = 12
    qkv_hidden_sizes = [192, 192, 96]
    mask_filter_value = -512.0

    dtype = "float32"
    input_array = np.random.random((batch_size, sequence_length, input_hidden_size)).astype(dtype)
    weight = np.random.normal(size=(input_hidden_size, sum(qkv_hidden_sizes))).astype(dtype) * 0.1
    bias = np.random.randn(sum(qkv_hidden_sizes)).astype(dtype)
    mask_index = np.random.randint(2, size=(batch_size, sequence_length)).astype("int32")
    relative_position_bias = np.random.randn(
        batch_size, num_heads, sequence_length, sequence_length
    ).astype(dtype)

    verify_attention(
        input_array,
        weight,
        bias,
        mask_index,
        num_heads,
        mask_filter_value,
        qkv_hidden_sizes,
        relative_position_bias,
    )


def _make_pad_expected_ir(input_shape, pads, mode="constant", value=0.0, opset=14, axes=None):
    len_dim = len(pads) // 2
    np_pads = [(pads[i], pads[i + len_dim]) for i in range(len_dim)]

    if axes is not None:
        rank = len(input_shape)
        full_pads = [(0, 0)] * rank
        for i, axis in enumerate(axes):
            axis = axis if axis >= 0 else axis + rank
            full_pads[axis] = np_pads[i]
        np_pads = full_pads

    if mode == "constant":
        out_shape = np.pad(
            np.empty(input_shape, dtype=np.float32),
            pad_width=np_pads,
            mode="constant",
            constant_values=value,
        ).shape
    else:
        out_shape = np.pad(
            np.empty(input_shape, dtype=np.float32), pad_width=np_pads, mode=mode
        ).shape
    input_shape = tuple(input_shape)
    out_shape = tuple(out_shape)
    pads_shape = (len(pads),)
    axes_shape = None if axes is None else (len(axes),)

    if mode == "constant" and opset >= 11:

        @I.ir_module
        class ExpectedPadConstantWithInputs:
            @T.prim_func(private=True, s_tir=True)
            def pad(input: T.handle, PadInput: T.handle):
                T.evaluate(0)

            @R.function
            def main(
                input: R.Tensor(input_shape, dtype="float32"),
                pads: R.Tensor(pads_shape, dtype="int64"),
                constant_value: R.Tensor((1,), dtype="float32"),
            ) -> R.Tensor(out_shape, dtype="float32"):
                R.func_attr({"num_input": 1})
                cls = ExpectedPadConstantWithInputs
                with R.dataflow():
                    lv = R.call_tir(
                        cls.pad,
                        (input,),
                        out_ty=R.Tensor(out_shape, dtype="float32"),
                    )
                    gv: R.Tensor(out_shape, dtype="float32") = lv
                    R.output(gv)
                return gv

        return ExpectedPadConstantWithInputs

    if mode == "constant":

        @I.ir_module
        class ExpectedPadConstantAttrs:
            @T.prim_func(private=True, s_tir=True)
            def pad(input: T.handle, PadInput: T.handle):
                T.evaluate(0)

            @R.function
            def main(input: R.Tensor(input_shape, dtype="float32")) -> R.Tensor(
                out_shape, dtype="float32"
            ):
                R.func_attr({"num_input": 1})
                cls = ExpectedPadConstantAttrs
                with R.dataflow():
                    lv = R.call_tir(
                        cls.pad,
                        (input,),
                        out_ty=R.Tensor(out_shape, dtype="float32"),
                    )
                    gv: R.Tensor(out_shape, dtype="float32") = lv
                    R.output(gv)
                return gv

        return ExpectedPadConstantAttrs

    if mode == "reflect" and opset >= 11:

        @I.ir_module
        class ExpectedPadReflectWithInputs:
            @T.prim_func(private=True, s_tir=True)
            def mirror_pad(input: T.handle, MirrorPadInput: T.handle):
                T.evaluate(0)

            @R.function
            def main(
                input: R.Tensor(input_shape, dtype="float32"),
                pads: R.Tensor(pads_shape, dtype="int64"),
            ) -> R.Tensor(out_shape, dtype="float32"):
                R.func_attr({"num_input": 1})
                cls = ExpectedPadReflectWithInputs
                with R.dataflow():
                    lv = R.call_tir(
                        cls.mirror_pad,
                        (input,),
                        out_ty=R.Tensor(out_shape, dtype="float32"),
                    )
                    gv: R.Tensor(out_shape, dtype="float32") = lv
                    R.output(gv)
                return gv

        return ExpectedPadReflectWithInputs

    if mode == "reflect":

        @I.ir_module
        class ExpectedPadReflectAttrs:
            @T.prim_func(private=True, s_tir=True)
            def mirror_pad(input: T.handle, MirrorPadInput: T.handle):
                T.evaluate(0)

            @R.function
            def main(input: R.Tensor(input_shape, dtype="float32")) -> R.Tensor(
                out_shape, dtype="float32"
            ):
                R.func_attr({"num_input": 1})
                cls = ExpectedPadReflectAttrs
                with R.dataflow():
                    lv = R.call_tir(
                        cls.mirror_pad,
                        (input,),
                        out_ty=R.Tensor(out_shape, dtype="float32"),
                    )
                    gv: R.Tensor(out_shape, dtype="float32") = lv
                    R.output(gv)
                return gv

        return ExpectedPadReflectAttrs

    if mode == "edge" and opset >= 11:

        @I.ir_module
        class ExpectedPadEdgeWithInputs:
            @T.prim_func(private=True, s_tir=True)
            def replicate_pad(input: T.handle, ReplicatePadInput: T.handle):
                T.evaluate(0)

            @R.function
            def main(
                input: R.Tensor(input_shape, dtype="float32"),
                pads: R.Tensor(pads_shape, dtype="int64"),
            ) -> R.Tensor(out_shape, dtype="float32"):
                R.func_attr({"num_input": 1})
                cls = ExpectedPadEdgeWithInputs
                with R.dataflow():
                    lv = R.call_tir(
                        cls.replicate_pad,
                        (input,),
                        out_ty=R.Tensor(out_shape, dtype="float32"),
                    )
                    gv: R.Tensor(out_shape, dtype="float32") = lv
                    R.output(gv)
                return gv

        return ExpectedPadEdgeWithInputs

    if mode == "edge":

        @I.ir_module
        class ExpectedPadEdgeAttrs:
            @T.prim_func(private=True, s_tir=True)
            def replicate_pad(input: T.handle, ReplicatePadInput: T.handle):
                T.evaluate(0)

            @R.function
            def main(input: R.Tensor(input_shape, dtype="float32")) -> R.Tensor(
                out_shape, dtype="float32"
            ):
                R.func_attr({"num_input": 1})
                cls = ExpectedPadEdgeAttrs
                with R.dataflow():
                    lv = R.call_tir(
                        cls.replicate_pad,
                        (input,),
                        out_ty=R.Tensor(out_shape, dtype="float32"),
                    )
                    gv: R.Tensor(out_shape, dtype="float32") = lv
                    R.output(gv)
                return gv

        return ExpectedPadEdgeAttrs

    if mode == "wrap" and opset >= 19:
        if axes is None:

            @I.ir_module
            class ExpectedPadWrapWithInputs:
                @T.prim_func(private=True, s_tir=True)
                def circular_pad(input: T.handle, CircularPadInput: T.handle):
                    T.evaluate(0)

                @R.function
                def main(
                    input: R.Tensor(input_shape, dtype="float32"),
                    pads: R.Tensor(pads_shape, dtype="int64"),
                ) -> R.Tensor(out_shape, dtype="float32"):
                    R.func_attr({"num_input": 1})
                    cls = ExpectedPadWrapWithInputs
                    with R.dataflow():
                        lv = R.call_tir(
                            cls.circular_pad,
                            (input,),
                            out_ty=R.Tensor(out_shape, dtype="float32"),
                        )
                        gv: R.Tensor(out_shape, dtype="float32") = lv
                        R.output(gv)
                    return gv

            return ExpectedPadWrapWithInputs

        @I.ir_module
        class ExpectedPadWrapWithAxes:
            @T.prim_func(private=True, s_tir=True)
            def circular_pad(input: T.handle, CircularPadInput: T.handle):
                T.evaluate(0)

            @R.function
            def main(
                input: R.Tensor(input_shape, dtype="float32"),
                pads: R.Tensor(pads_shape, dtype="int64"),
                axes: R.Tensor(axes_shape, dtype="int64"),
            ) -> R.Tensor(out_shape, dtype="float32"):
                R.func_attr({"num_input": 1})
                cls = ExpectedPadWrapWithAxes
                with R.dataflow():
                    lv = R.call_tir(
                        cls.circular_pad,
                        (input,),
                        out_ty=R.Tensor(out_shape, dtype="float32"),
                    )
                    gv: R.Tensor(out_shape, dtype="float32") = lv
                    R.output(gv)
                return gv

        return ExpectedPadWrapWithAxes

    raise AssertionError(f"No Pad expected IR for mode={mode}, opset={opset}")


@pytest.mark.parametrize("dynamic", [True, False])
def test_pad(dynamic):
    if dynamic:
        pytest.skip("Dynamic pad not supported")

    def verify_pad(input_shape, pads, expected, mode="constant", value=0.0, opset=14, axes=None):
        len_dim = len(pads) // 2
        np_pads = [(pads[i], pads[i + len_dim]) for i in range(len_dim)]

        if axes is not None:
            rank = len(input_shape)
            full_pads = [(0, 0)] * rank
            for i, axis in enumerate(axes):
                axis = axis if axis >= 0 else axis + rank
                full_pads[axis] = np_pads[i]
            np_pads = full_pads

        pads = np.array(pads, dtype=np.int64)
        #  onnx graph
        if mode in ["edge", "reflect", "wrap"]:
            outdata = np.pad(np.empty(input_shape, dtype=np.float32), pad_width=np_pads, mode=mode)

            node_inputs = ["input", "pads"]
            initializer = [helper.make_tensor("pads", TensorProto.INT64, (len(pads),), pads)]

            if axes is not None:
                axes = np.array(axes, dtype=np.int64)
                node_inputs = ["input", "pads", "", "axes"]
                initializer.append(
                    helper.make_tensor("axes", TensorProto.INT64, (len(axes),), axes)
                )

            node = helper.make_node("Pad", inputs=node_inputs, outputs=["output"], mode=mode)
            graph = helper.make_graph(
                [node],
                "pad_test",
                inputs=[
                    helper.make_tensor_value_info("input", TensorProto.FLOAT, list(input_shape))
                ],
                initializer=initializer,
                outputs=[
                    helper.make_tensor_value_info("output", TensorProto.FLOAT, list(outdata.shape))
                ],
            )
        else:
            outdata = np.pad(
                np.empty(input_shape, dtype=np.float32),
                pad_width=np_pads,
                mode="constant",
                constant_values=value,
            )
            node = helper.make_node(
                "Pad",
                inputs=["input", "pads", "constant_value"],
                outputs=["output"],
                mode="constant",
            )
            graph = helper.make_graph(
                [node],
                "pad_test",
                inputs=[
                    helper.make_tensor_value_info("input", TensorProto.FLOAT, list(input_shape))
                ],
                initializer=[
                    helper.make_tensor("pads", TensorProto.INT64, (len(pads),), pads),
                    helper.make_tensor("constant_value", TensorProto.FLOAT, (1,), [value]),
                ],
                outputs=[
                    helper.make_tensor_value_info("output", TensorProto.FLOAT, list(outdata.shape))
                ],
            )
        model = helper.make_model(graph, producer_name="pad_test")
        model.opset_import[0].version = 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, pads, mode, value, opset, axes in [
        ((2, 2), [0, 1, 0, 0], "constant", 0.0, 14, None),
        ((2, 3), [1, 0, 0, 1], "constant", 0.0, 14, None),
        ((3, 2), [0, 0, 1, 0], "constant", 5.0, 14, None),
        ((1, 3, 4, 5), [0, 1, 1, 1, 0, 0, 1, 1], "reflect", 0.0, 14, None),
        ((2, 3), [1, 1, 1, 1], "edge", 0.0, 14, None),
        ((1, 3, 4, 5), [0, 1, 1, 1, 0, 0, 1, 1], "edge", 0.0, 14, None),
        ((1, 3, 4), [0, 0, 2, 0, 0, 2], "wrap", 0.0, 19, None),
        ((1, 3, 4), [2, 2], "wrap", 0.0, 19, [2]),
        ((1, 3, 4), [1, 2, 1, 2], "wrap", 0.0, 19, [1, 2]),
    ]:
        verify_pad(
            input_shape,
            pads,
            _make_pad_expected_ir(
                input_shape, pads, mode=mode, value=value, opset=opset, axes=axes
            ),
            mode,
            value,
            opset,
            axes,
        )


@pytest.mark.parametrize("dynamic", [True, False])
def test_pad_v2(dynamic):
    if dynamic:
        pytest.skip("Dynamic pad not supported")

    def verify_pad(input_shape, pads, expected, mode="constant", value=0.0):
        len_dim = len(pads) // 2
        np_pads = [(pads[i], pads[i + len_dim]) for i in range(len_dim)]
        pads = np.array(pads)
        #  onnx graph
        if mode in ["edge", "reflect"]:
            outdata = np.pad(np.empty(input_shape, dtype=np.float32), pad_width=np_pads, mode=mode)
            node = helper.make_node(
                "Pad", inputs=["input"], outputs=["output"], mode=mode, pads=pads
            )
            graph = helper.make_graph(
                [node],
                "pad_test",
                inputs=[
                    helper.make_tensor_value_info("input", TensorProto.FLOAT, list(input_shape))
                ],
                outputs=[
                    helper.make_tensor_value_info("output", TensorProto.FLOAT, list(outdata.shape))
                ],
            )
        else:
            outdata = np.pad(
                np.empty(input_shape, dtype=np.float32),
                pad_width=np_pads,
                mode="constant",
                constant_values=value,
            )
            node = helper.make_node(
                "Pad",
                inputs=["input"],
                outputs=["output"],
                mode="constant",
                pads=pads,
                value=value,
            )
            graph = helper.make_graph(
                [node],
                "pad_test",
                inputs=[
                    helper.make_tensor_value_info("input", TensorProto.FLOAT, list(input_shape))
                ],
                outputs=[
                    helper.make_tensor_value_info("output", TensorProto.FLOAT, list(outdata.shape))
                ],
            )
        model = helper.make_model(graph, producer_name="pad_test")
        model.opset_import[0].version = 10
        tvm_model = from_onnx(model, opset=10, keep_params_in_input=True)
        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, pads, mode, value in [
        ((2, 2), [0, 1, 0, 0], "constant", 0.0),
        ((2, 3), [1, 0, 0, 1], "constant", 0.0),
        ((3, 2), [0, 0, 1, 0], "constant", 5.0),
        ((1, 3, 4, 5), [0, 1, 1, 1, 0, 0, 1, 1], "reflect", 0.0),
        ((2, 3), [1, 1, 1, 1], "edge", 0.0),
        ((1, 3, 4, 5), [0, 1, 1, 1, 0, 0, 1, 1], "edge", 0.0),
    ]:
        verify_pad(
            input_shape,
            pads,
            _make_pad_expected_ir(input_shape, pads, mode=mode, value=value, opset=10),
            mode,
            value,
        )


def test_split():
    def verify_split(
        fp_arith,
        dynamic,
        indata_shape,
        outdata_shapes,
        split,
        expected,
        axis=0,
        pass_split=True,
        opset=11,
    ):
        indata = np.random.normal(size=indata_shape).astype(fp_arith)
        input_names = ["input"]
        initializer = []

        if split:
            split_index = range(len(split))
        else:
            split_index = range(len(outdata_shapes))

        indata_shape = list(indata.shape)
        if dynamic:
            indata_shape = ["?" for _ in range(len(indata.shape))]
            outdata_shapes = [["?" for _ in range(len(o))] for o in outdata_shapes]

        inputs = [
            helper.make_tensor_value_info(
                "input", helper.np_dtype_to_tensor_dtype(indata.dtype), indata_shape
            )
        ]

        split_constant = None
        if pass_split:
            if opset >= 13:
                np_split = np.array(split).astype(np.int64)
                split_constant = make_constant_node(
                    "split", onnx.TensorProto.INT64, list(np_split.shape), np_split
                )
                input_names.append("split")

        node = helper.make_node(
            "Split",
            inputs=input_names,
            outputs=[f"output_{i}" for i in range(len(split_index))],
            axis=axis,
        )

        if pass_split and opset < 13:
            split_attr = helper.make_attribute("split", split)
            node.attribute.append(split_attr)

        nodes = [split_constant, node] if split_constant else [node]

        graph = helper.make_graph(
            nodes,
            "split_test",
            inputs=inputs,
            initializer=initializer,
            outputs=[
                helper.make_tensor_value_info(
                    f"output_{i}",
                    helper.np_dtype_to_tensor_dtype(indata.dtype),
                    list(outdata_shapes[i]),
                )
                for i in range(len(split_index))
            ],
        )
        model = helper.make_model(graph, producer_name="split_test")
        tvm_model = from_onnx(model, opset=opset, keep_params_in_input=True)
        tvm.ir.assert_structural_equal(tvm_model, expected)

    def make_expected(fp_arith, dynamic, indata_shape, outdata_shapes, split, axis, pass_split):
        def shape_tuple(shape):
            if isinstance(shape, int):
                shape = (shape,)
            return tuple(shape)

        def expected_input_shape(shape):
            shape = shape_tuple(shape)
            if not dynamic:
                return shape
            return tuple(f"split_input_dim_{i}" for i in range(len(shape)))

        dtype = np.dtype(fp_arith).name
        input_shape = expected_input_shape(indata_shape)
        if not pass_split:
            indices_or_sections = len(outdata_shapes)
        elif len(outdata_shapes) == 1:
            indices_or_sections = 1
        else:
            indices_or_sections = list(np.cumsum(split)[:-1])

        if len(outdata_shapes) == 1:

            @I.ir_module
            class ExpectedSplitSingle:
                @R.function
                def main(input: R.Tensor(input_shape, dtype=dtype)):
                    R.func_attr({"num_input": 1})
                    with R.dataflow():
                        gv = R.split(input, indices_or_sections=indices_or_sections, axis=axis)
                        R.output(gv)
                    return gv

            return ExpectedSplitSingle

        if len(outdata_shapes) == 2:

            @I.ir_module
            class ExpectedSplitPair:
                @R.function
                def main(input: R.Tensor(input_shape, dtype=dtype)):
                    R.func_attr({"num_input": 1})
                    with R.dataflow():
                        lv = R.split(input, indices_or_sections=indices_or_sections, axis=axis)
                        lv1 = lv[0]
                        lv2 = lv[1]
                        gv = (lv1, lv2)
                        R.output(gv)
                    return gv

            return ExpectedSplitPair

        assert len(outdata_shapes) == 3

        @I.ir_module
        class ExpectedSplitTriple:
            @R.function
            def main(input: R.Tensor(input_shape, dtype=dtype)):
                R.func_attr({"num_input": 1})
                with R.dataflow():
                    lv = R.split(input, indices_or_sections=indices_or_sections, axis=axis)
                    lv1 = lv[0]
                    lv2 = lv[1]
                    lv3 = lv[2]
                    gv = (lv1, lv2, lv3)
                    R.output(gv)
                return gv

        return ExpectedSplitTriple

    split_cases = [
        (6, [[2], [2], [2]], [2, 2, 2], 0, True, 11),
        (6, [[2], [2], [2]], [2, 2, 2], 0, False, 11),
        (6, [[2], [1], [3]], [2, 1, 3], 0, True, 11),
        (6, [[2], [1], [3]], [2, 1, 3], 0, True, 13),
        ((4, 4), [[2, 2], [2, 2]], [2, 2], 1, True, 11),
        ((4, 4), [[2, 2], [2, 2]], [2, 2], 1, True, 13),
        (3, [[1], [1], [1]], False, 0, False, 11),
        (1, [[1]], [1], 0, True, 11),
        ((1, 2), [[2]], [2], 1, True, 11),
        ((1, 2), [[2]], [1], 0, True, 11),
    ]

    for fp_arith in [np.float16, np.float32]:
        for dynamic in [True, False]:
            for indata_shape, outdata_shapes, split, axis, pass_split, opset in split_cases:
                verify_split(
                    fp_arith,
                    dynamic,
                    indata_shape,
                    outdata_shapes,
                    split,
                    make_expected(
                        fp_arith, dynamic, indata_shape, outdata_shapes, split, axis, pass_split
                    ),
                    axis=axis,
                    pass_split=pass_split,
                    opset=opset,
                )


def test_tile():
    def verify_tile(dynamic, in_shape, repeats, out_shape, expected):
        node = helper.make_node("Tile", inputs=["input", "repeats"], outputs=["out"])

        model_in_shape = list(in_shape)
        model_out_shape = list(out_shape)
        if dynamic:
            model_in_shape = ["?" for _ in range(len(in_shape))]
            model_out_shape = ["?" for _ in range(len(out_shape))]

        graph = helper.make_graph(
            [node],
            "tile_test",
            inputs=[
                helper.make_tensor_value_info("input", TensorProto.FLOAT, model_in_shape),
            ],
            initializer=[
                helper.make_tensor("repeats", TensorProto.INT64, list(repeats.shape), repeats)
            ],
            outputs=[helper.make_tensor_value_info("out", TensorProto.FLOAT, model_out_shape)],
        )

        model = helper.make_model(
            graph, producer_name="tile_test", opset_imports=[helper.make_opsetid("", 14)]
        )
        tvm_model = from_onnx(model, keep_params_in_input=True)
        assert len(tvm_model["main"].attrs["params"]) == 1
        tvm_model["main"] = tvm_model["main"].without_attr("params")
        expected = tvm.IRModule(expected.functions)
        expected.update_func(expected.get_global_var("tile"), tvm_model["tile"])
        tvm.ir.assert_structural_equal(tvm_model, expected)

    @I.ir_module
    class ExpectedTileDynamicInput:
        @T.prim_func(private=True, s_tir=True)
        def tile(input: T.handle, T_tile: T.handle):
            T.evaluate(0)

        @R.function
        def main(
            input: R.Tensor(
                (
                    "tile_input_dim_0",
                    "tile_input_dim_1",
                    "tile_input_dim_2",
                    "tile_input_dim_3",
                ),
                dtype="float32",
            ),
            repeats: R.Tensor((4,), dtype="int64"),
        ) -> R.Tensor(
            (
                "tile_input_dim_0 * 2",
                "tile_input_dim_1",
                "tile_input_dim_2 * 3",
                "tile_input_dim_3 * 2",
            ),
            dtype="float32",
        ):
            tile_input_dim_0 = T.int64()
            tile_input_dim_1 = T.int64()
            tile_input_dim_2 = T.int64()
            tile_input_dim_3 = T.int64()
            R.func_attr({"num_input": 1})
            cls = ExpectedTileDynamicInput
            with R.dataflow():
                lv = R.call_tir(
                    cls.tile,
                    (input,),
                    out_ty=R.Tensor(
                        (
                            tile_input_dim_0 * 2,
                            tile_input_dim_1,
                            tile_input_dim_2 * 3,
                            tile_input_dim_3 * 2,
                        ),
                        dtype="float32",
                    ),
                )
                gv: R.Tensor(
                    (
                        tile_input_dim_0 * 2,
                        tile_input_dim_1,
                        tile_input_dim_2 * 3,
                        tile_input_dim_3 * 2,
                    ),
                    dtype="float32",
                ) = lv
                R.output(gv)
            return gv

    @I.ir_module
    class ExpectedTileStaticInput:
        @T.prim_func(private=True, s_tir=True)
        def tile(input: T.handle, T_tile: T.handle):
            T.evaluate(0)

        @R.function
        def main(
            input: R.Tensor((2, 3, 4, 5), dtype="float32"),
            repeats: R.Tensor((4,), dtype="int64"),
        ) -> R.Tensor((4, 3, 12, 10), dtype="float32"):
            R.func_attr({"num_input": 1})
            cls = ExpectedTileStaticInput
            with R.dataflow():
                lv = R.call_tir(
                    cls.tile,
                    (input,),
                    out_ty=R.Tensor((4, 3, 12, 10), dtype="float32"),
                )
                gv: R.Tensor((4, 3, 12, 10), dtype="float32") = lv
                R.output(gv)
            return gv

    x = np.random.rand(2, 3, 4, 5).astype(np.float32)
    repeats = np.array([2, 1, 3, 2], dtype=np.int64)
    z_array = np.tile(x, repeats)
    verify_tile(True, x.shape, repeats, z_array.shape, ExpectedTileDynamicInput)
    verify_tile(False, x.shape, repeats, z_array.shape, ExpectedTileStaticInput)


def test_tile_dynamic_repeats():
    def verify_tile_dynamic_repeats(dynamic_input, in_shape, repeats, expected):
        out_shape = np.tile(np.empty(in_shape, dtype=np.float32), repeats).shape

        input_shape = ["?" for _ in in_shape] if dynamic_input else list(in_shape)
        output_shape = ["?" for _ in out_shape] if dynamic_input else list(out_shape)

        node = helper.make_node("Tile", inputs=["input", "repeats"], outputs=["out"])
        graph = helper.make_graph(
            [node],
            "tile_dynamic_repeats_test",
            inputs=[
                helper.make_tensor_value_info("input", TensorProto.FLOAT, input_shape),
                helper.make_tensor_value_info("repeats", TensorProto.INT64, [len(repeats)]),
            ],
            outputs=[helper.make_tensor_value_info("out", TensorProto.FLOAT, output_shape)],
        )
        model = helper.make_model(
            graph,
            producer_name="tile_dynamic_repeats_test",
            opset_imports=[helper.make_opsetid("", 13)],
        )

        tvm_model = from_onnx(model, opset=13, keep_params_in_input=True)
        expected = tvm.IRModule(expected.functions)
        expected.update_func(expected.get_global_var("dyn_tile"), tvm_model["dyn_tile"])
        tvm.ir.assert_structural_equal(tvm_model, expected)

    def make_expected(dynamic_input, in_shape):
        rank = len(in_shape)
        input_shape = (
            tuple(f"tile_data_dim_{i}" for i in range(rank)) if dynamic_input else tuple(in_shape)
        )

        if rank == 2:

            @I.ir_module
            class ExpectedTileRank2:
                @T.prim_func(private=True, s_tir=True)
                def dyn_tile(input: T.handle, var_T_tile: T.handle):
                    T.evaluate(0)

                @R.function
                def main(
                    input: R.Tensor(input_shape, dtype="float32"),
                    repeats: R.Tensor((2,), dtype="int64"),
                ) -> R.Tensor(dtype="float32", ndim=2):
                    tile_dim_0 = T.int64()
                    tile_dim_1 = T.int64()
                    R.func_attr({"num_input": 2})
                    cls = ExpectedTileRank2
                    with R.dataflow():
                        lv = R.shape_of(input)
                        lv1: R.Tensor((2,), dtype="int64") = R.shape_to_tensor(lv)
                        lv2: R.Tensor((2,), dtype="int64") = R.multiply(repeats, lv1)
                        lv3: R.Shape([tile_dim_0, tile_dim_1]) = R.match_cast(
                            R.tensor_to_shape(lv2), R.Shape([tile_dim_0, tile_dim_1])
                        )
                        lv4 = R.call_tir(
                            cls.dyn_tile,
                            (input,),
                            out_ty=R.Tensor((tile_dim_0, tile_dim_1), dtype="float32"),
                        )
                        gv: R.Tensor((tile_dim_0, tile_dim_1), dtype="float32") = lv4
                        R.output(gv)
                    return gv

            return ExpectedTileRank2

        if rank == 3:

            @I.ir_module
            class ExpectedTileRank3:
                @T.prim_func(private=True, s_tir=True)
                def dyn_tile(input: T.handle, var_T_tile: T.handle):
                    T.evaluate(0)

                @R.function
                def main(
                    input: R.Tensor(input_shape, dtype="float32"),
                    repeats: R.Tensor((3,), dtype="int64"),
                ) -> R.Tensor(dtype="float32", ndim=3):
                    tile_dim_0 = T.int64()
                    tile_dim_1 = T.int64()
                    tile_dim_2 = T.int64()
                    R.func_attr({"num_input": 2})
                    cls = ExpectedTileRank3
                    with R.dataflow():
                        lv = R.shape_of(input)
                        lv1: R.Tensor((3,), dtype="int64") = R.shape_to_tensor(lv)
                        lv2: R.Tensor((3,), dtype="int64") = R.multiply(repeats, lv1)
                        lv3: R.Shape([tile_dim_0, tile_dim_1, tile_dim_2]) = R.match_cast(
                            R.tensor_to_shape(lv2),
                            R.Shape([tile_dim_0, tile_dim_1, tile_dim_2]),
                        )
                        lv4 = R.call_tir(
                            cls.dyn_tile,
                            (input,),
                            out_ty=R.Tensor((tile_dim_0, tile_dim_1, tile_dim_2), dtype="float32"),
                        )
                        gv: R.Tensor((tile_dim_0, tile_dim_1, tile_dim_2), dtype="float32") = lv4
                        R.output(gv)
                    return gv

            return ExpectedTileRank3

        if rank == 4:

            @I.ir_module
            class ExpectedTileRank4:
                @T.prim_func(private=True, s_tir=True)
                def dyn_tile(input: T.handle, var_T_tile: T.handle):
                    T.evaluate(0)

                @R.function
                def main(
                    input: R.Tensor(input_shape, dtype="float32"),
                    repeats: R.Tensor((4,), dtype="int64"),
                ) -> R.Tensor(dtype="float32", ndim=4):
                    tile_dim_0 = T.int64()
                    tile_dim_1 = T.int64()
                    tile_dim_2 = T.int64()
                    tile_dim_3 = T.int64()
                    R.func_attr({"num_input": 2})
                    cls = ExpectedTileRank4
                    with R.dataflow():
                        lv = R.shape_of(input)
                        lv1: R.Tensor((4,), dtype="int64") = R.shape_to_tensor(lv)
                        lv2: R.Tensor((4,), dtype="int64") = R.multiply(repeats, lv1)
                        lv3: R.Shape([tile_dim_0, tile_dim_1, tile_dim_2, tile_dim_3]) = (
                            R.match_cast(
                                R.tensor_to_shape(lv2),
                                R.Shape([tile_dim_0, tile_dim_1, tile_dim_2, tile_dim_3]),
                            )
                        )
                        lv4 = R.call_tir(
                            cls.dyn_tile,
                            (input,),
                            out_ty=R.Tensor(
                                (tile_dim_0, tile_dim_1, tile_dim_2, tile_dim_3),
                                dtype="float32",
                            ),
                        )
                        gv: R.Tensor(
                            (tile_dim_0, tile_dim_1, tile_dim_2, tile_dim_3), dtype="float32"
                        ) = lv4
                        R.output(gv)
                    return gv

            return ExpectedTileRank4

        raise AssertionError(f"No dynamic Tile expected IR for rank {rank}")

    tile_cases = [
        (True, (2, 3), np.array([2, 2], dtype=np.int64)),
        (True, (2, 3, 4), np.array([2, 2, 1], dtype=np.int64)),
        (True, (2, 3, 4, 5), np.array([1, 2, 1, 2], dtype=np.int64)),
        (False, (2, 3), np.array([2, 2], dtype=np.int64)),
        (False, (2, 3, 4), np.array([2, 2, 1], dtype=np.int64)),
        (False, (2, 3, 4, 5), np.array([1, 2, 1, 2], dtype=np.int64)),
    ]
    for dynamic_input, in_shape, repeats in tile_cases:
        verify_tile_dynamic_repeats(
            dynamic_input, in_shape, repeats, make_expected(dynamic_input, in_shape)
        )


def _generate_roi_cases():
    # Base case when with_roi is False
    roi_list = [
        pytest.param(False, None, False, id="no_roi"),
    ]

    # Valid when with_roi is True and with_constant is True/False
    roi_cases = [
        [],
        [0.0, 0.0, 0.0, 0.0],
        [0.0, 0.0, 1.0, 1.0],
        [0.1, 0.1, 0.9, 0.9],
        [0.2, 0.2, 0.8, 0.8],
        [0.3, 0.3, 0.7, 0.7],
        [0.4, 0.4, 0.6, 0.6],
        [0.5, 0.5, 0.5, 0.5],
        [0.1, 0.2, 0.9, 0.8],
    ]
    for roi in roi_cases:
        roi_list.append(pytest.param(True, roi, True, id=f"roi_{'_'.join(str(x) for x in roi)}"))
        roi_list.append(pytest.param(True, roi, False, id=f"roi_{'_'.join(str(x) for x in roi)}"))

    return roi_list


@pytest.mark.parametrize("with_roi, roi_list, with_constant", _generate_roi_cases())
def test_resize(with_roi, roi_list, with_constant):
    nodes = []
    resize_node = helper.make_node(
        "Resize", ["X", "roi" if with_roi else "", "scales"], ["Y"], mode="cubic"
    )

    if with_roi and with_constant:
        roi_tensor = helper.make_tensor(
            name="roi",
            data_type=TensorProto.FLOAT,
            dims=[len(roi_list)],
            vals=roi_list,
        )

        roi_const_node = helper.make_node(
            "Constant",
            inputs=[],
            outputs=["roi"],
            value=roi_tensor,
        )
        nodes.append(roi_const_node)

    nodes.append(resize_node)

    initializers = [
        helper.make_tensor("scales", TensorProto.FLOAT, [4], [1.0, 1.0, 2.0, 2.0]),
    ]

    if with_roi and not with_constant:
        initializers.append(helper.make_tensor("roi", TensorProto.FLOAT, [len(roi_list)], roi_list))

    graph = helper.make_graph(
        nodes,
        "resize_test",
        inputs=[
            helper.make_tensor_value_info("X", TensorProto.FLOAT, [1, 3, 32, 32]),
        ],
        initializer=initializers,
        outputs=[
            helper.make_tensor_value_info("Y", TensorProto.FLOAT, [1, 3, 64, 64]),
        ],
    )

    model = helper.make_model(graph, producer_name="resize_test")
    check_correctness(model)


def test_resize_dynamic_roi_tf_crop_and_resize():
    """ROI is a graph input (not initializer), lowered through TOPI dynamic-ROI path."""
    resize_node = helper.make_node(
        "Resize",
        ["X", "roi", "scales"],
        ["Y"],
        mode="linear",
        coordinate_transformation_mode="tf_crop_and_resize",
    )
    graph = helper.make_graph(
        [resize_node],
        "resize_dynamic_roi",
        inputs=[
            helper.make_tensor_value_info("X", TensorProto.FLOAT, [1, 3, 32, 32]),
            helper.make_tensor_value_info("roi", TensorProto.FLOAT, [8]),
        ],
        initializer=[
            helper.make_tensor("scales", TensorProto.FLOAT, [4], [1.0, 1.0, 2.0, 2.0]),
        ],
        outputs=[
            helper.make_tensor_value_info("Y", TensorProto.FLOAT, [1, 3, 64, 64]),
        ],
    )
    model = helper.make_model(graph, producer_name="resize_dynamic_roi")
    tvm_model = from_onnx(model, keep_params_in_input=True)
    seen_call_tir = False

    def _visit(expr):
        nonlocal seen_call_tir
        if isinstance(expr, relax.Call) and isinstance(expr.op, tvm.ir.Op):
            if expr.op.name == "relax.call_tir":
                seen_call_tir = True

    relax.analysis.post_order_visit(tvm_model["main"].body, _visit)
    assert seen_call_tir


def test_resize_dynamic_roi_3d_tf_crop_and_resize():
    """5-D NCDHW: ROI is a graph input; covers dynamic-ROI TOPI resize3d path."""
    resize_node = helper.make_node(
        "Resize",
        ["X", "roi", "scales"],
        ["Y"],
        mode="linear",
        coordinate_transformation_mode="tf_crop_and_resize",
    )
    graph = helper.make_graph(
        [resize_node],
        "resize_dynamic_roi_3d",
        inputs=[
            helper.make_tensor_value_info("X", TensorProto.FLOAT, [1, 1, 3, 4, 5]),
            helper.make_tensor_value_info("roi", TensorProto.FLOAT, [10]),
        ],
        initializer=[
            helper.make_tensor("scales", TensorProto.FLOAT, [5], [1.0, 1.0, 2.0, 2.0, 2.0]),
        ],
        outputs=[
            helper.make_tensor_value_info("Y", TensorProto.FLOAT, [1, 1, 6, 8, 10]),
        ],
    )
    model = helper.make_model(
        graph,
        producer_name="resize_dynamic_roi_3d",
        opset_imports=[helper.make_opsetid("", 18)],
    )
    tvm_model = from_onnx(model, opset=18, keep_params_in_input=True)
    seen_call_tir = False

    def _visit(expr):
        nonlocal seen_call_tir
        if isinstance(expr, relax.Call) and isinstance(expr.op, tvm.ir.Op):
            if expr.op.name == "relax.call_tir":
                seen_call_tir = True

    relax.analysis.post_order_visit(tvm_model["main"].body, _visit)
    assert seen_call_tir


def test_resize_nd_sizes():
    cases = [
        ("resize1d", [1, 1, 4], [1, 1, 7]),
        ("resize2d", [1, 1, 4, 5], [1, 1, 6, 7]),
        ("resize3d", [1, 1, 3, 4, 5], [1, 1, 4, 6, 7]),
    ]

    for name, input_shape, sizes in cases:
        resize_node = helper.make_node(
            "Resize",
            ["X", "", "", "sizes"],
            ["Y"],
            mode="nearest",
            coordinate_transformation_mode="asymmetric",
            nearest_mode="floor",
        )

        graph = helper.make_graph(
            [resize_node],
            name,
            inputs=[
                helper.make_tensor_value_info("X", TensorProto.FLOAT, input_shape),
            ],
            initializer=[
                helper.make_tensor("sizes", TensorProto.INT64, [len(sizes)], sizes),
            ],
            outputs=[
                helper.make_tensor_value_info("Y", TensorProto.FLOAT, sizes),
            ],
        )

        model = helper.make_model(
            graph, producer_name=name, opset_imports=[helper.make_opsetid("", 18)]
        )
        if name != "resize1d":
            check_correctness(model, opset=18)
            continue

        tvm_model = from_onnx(model, opset=18, keep_params_in_input=True)
        seen_call_tir = False

        def _visit(expr):
            nonlocal seen_call_tir
            if isinstance(expr, relax.Call) and isinstance(expr.op, tvm.ir.Op):
                if expr.op.name == "relax.call_tir":
                    seen_call_tir = True

        relax.analysis.post_order_visit(tvm_model["main"].body, _visit)
        assert seen_call_tir


def test_resize_5d_emits_relax_resize3d():
    resize_node = helper.make_node(
        "Resize",
        ["X", "", "", "sizes"],
        ["Y"],
        mode="nearest",
        coordinate_transformation_mode="asymmetric",
        nearest_mode="floor",
    )
    graph = helper.make_graph(
        [resize_node],
        "resize3d_ir_check",
        inputs=[helper.make_tensor_value_info("X", TensorProto.FLOAT, [1, 1, 3, 4, 5])],
        initializer=[helper.make_tensor("sizes", TensorProto.INT64, [5], [1, 1, 4, 6, 7])],
        outputs=[helper.make_tensor_value_info("Y", TensorProto.FLOAT, [1, 1, 4, 6, 7])],
    )
    model = helper.make_model(graph, producer_name="resize3d_ir_check")
    tvm_model = from_onnx(model, opset=18, keep_params_in_input=True)

    seen_resize3d = False

    def _visit(expr):
        nonlocal seen_resize3d
        if isinstance(expr, relax.Call) and isinstance(expr.op, tvm.ir.Op):
            if expr.op.name == "relax.image.resize3d":
                seen_resize3d = True

    relax.analysis.post_order_visit(tvm_model["main"].body, _visit)
    assert seen_resize3d


def test_einsum():
    eqn = "ij->i"
    einsum_node = helper.make_node("Einsum", ["x"], ["y"], equation=eqn)

    graph = helper.make_graph(
        [einsum_node],
        "einsum_test",
        inputs=[
            helper.make_tensor_value_info("x", TensorProto.FLOAT, [3, 4]),
        ],
        outputs=[
            helper.make_tensor_value_info("y", TensorProto.FLOAT, [3]),
        ],
    )

    model = helper.make_model(graph, producer_name="einsum_test")
    tvm_model = from_onnx(model, keep_params_in_input=True)

    @I.ir_module
    class Expected:
        @T.prim_func(private=True, s_tir=True)
        def einsum(x: T.handle, T_einsum: T.handle):
            T.evaluate(0)

        @R.function
        def main(x: R.Tensor((3, 4), dtype="float32")) -> R.Tensor((3,), dtype="float32"):
            R.func_attr({"num_input": 1})
            cls = Expected
            with R.dataflow():
                lv = R.call_tir(cls.einsum, (x,), out_ty=R.Tensor((3,), dtype="float32"))
                gv: R.Tensor((3,), dtype="float32") = lv
                R.output(gv)
            return gv

    expected = tvm.IRModule(Expected.functions)
    expected.update_func(expected.get_global_var("einsum"), tvm_model["einsum"])
    tvm.ir.assert_structural_equal(tvm_model, expected)


def test_range():
    range_node = helper.make_node(
        "Range",
        ["start", "limit", "delta"],
        ["output"],
    )

    graph = helper.make_graph(
        [range_node],
        "range_test",
        inputs=[],
        initializer=[
            helper.make_tensor("start", TensorProto.INT64, [], [1]),
            helper.make_tensor("limit", TensorProto.INT64, [], [5]),
            helper.make_tensor("delta", TensorProto.INT64, [], [2]),
        ],
        outputs=[
            helper.make_tensor_value_info("output", TensorProto.INT64, [2]),
        ],
    )

    model = helper.make_model(graph, producer_name="range_test")
    tvm_model = from_onnx(model, keep_params_in_input=True)
    tvm_model["main"] = tvm_model["main"].without_attr("params")

    @I.ir_module
    class Expected:
        @R.function
        def main(
            start: R.Tensor((), dtype="int64"),
            limit: R.Tensor((), dtype="int64"),
            delta: R.Tensor((), dtype="int64"),
        ) -> R.Tensor((2,), dtype="int64"):
            R.func_attr({"num_input": 0})
            with R.dataflow():
                gv: R.Tensor((2,), dtype="int64") = R.const(
                    np.array([1, 3], dtype=np.int64), "int64"
                )
                R.output(gv)
            return gv

    tvm.ir.assert_structural_equal(tvm_model, Expected)


def test_batch_norm():
    batch_norm_node = helper.make_node(
        "BatchNormalization", ["x", "s", "bias", "mean", "var"], ["y"], epsilon=1e-2
    )
    graph = helper.make_graph(
        [batch_norm_node],
        "batch_norm_test",
        inputs=[
            helper.make_tensor_value_info("x", TensorProto.FLOAT, [2, 3, 4, 5]),
            helper.make_tensor_value_info("s", TensorProto.FLOAT, [3]),
            helper.make_tensor_value_info("bias", TensorProto.FLOAT, [3]),
            helper.make_tensor_value_info("mean", TensorProto.FLOAT, [3]),
            helper.make_tensor_value_info("var", TensorProto.FLOAT, [3]),
        ],
        outputs=[helper.make_tensor_value_info("y", TensorProto.FLOAT, [2, 3, 4, 5])],
    )

    model = helper.make_model(graph, producer_name="batch_norm_test")
    check_correctness(model, opset=15)


def test_batch_norm_defaults_to_inference_mode():
    batch_norm_node = helper.make_node(
        "BatchNormalization", ["x", "s", "bias", "mean", "var"], ["y"], epsilon=1e-2
    )
    graph = helper.make_graph(
        [batch_norm_node],
        "batch_norm_inference_attr_test",
        inputs=[
            helper.make_tensor_value_info("x", TensorProto.FLOAT, [2, 3, 4, 5]),
            helper.make_tensor_value_info("s", TensorProto.FLOAT, [3]),
            helper.make_tensor_value_info("bias", TensorProto.FLOAT, [3]),
            helper.make_tensor_value_info("mean", TensorProto.FLOAT, [3]),
            helper.make_tensor_value_info("var", TensorProto.FLOAT, [3]),
        ],
        outputs=[helper.make_tensor_value_info("y", TensorProto.FLOAT, [2, 3, 4, 5])],
    )
    model = helper.make_model(graph, producer_name="batch_norm_inference_attr_test")
    model.opset_import[0].version = 15

    tvm_model = from_onnx(model, opset=15, keep_params_in_input=True)
    batch_norm_attrs = []

    def visit(expr):
        if isinstance(expr, relax.Call) and expr.op == tvm.ir.Op.get("relax.nn.batch_norm"):
            batch_norm_attrs.append(expr.attrs)

    relax.analysis.post_order_visit(tvm_model["main"], visit)

    assert len(batch_norm_attrs) == 1
    assert batch_norm_attrs[0].training is False


def test_batch_norm_mixed_dtype_params():
    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"),
    ]
    batch_norm_node = helper.make_node(
        "BatchNormalization",
        ["data", "gamma", "beta", "mean", "var"],
        ["output"],
        epsilon=1e-5,
        momentum=0.9,
        training_mode=0,
    )
    graph = helper.make_graph(
        [batch_norm_node],
        "mixed_dtype_batchnorm",
        [data],
        [output],
        initializer=params,
    )
    model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 15)])

    tvm_model = from_onnx(model, keep_params_in_input=False)

    assert tuple(dim.value for dim in tvm_model["main"].ret_ty.shape.values) == (1, 3, 2, 2)
    assert tvm_model["main"].ret_ty.dtype == "float16"

    batch_norm_calls = []

    def visit(expr):
        if isinstance(expr, relax.Call) and expr.op == tvm.ir.Op.get("relax.nn.batch_norm"):
            batch_norm_calls.append(expr)

    relax.analysis.post_order_visit(tvm_model["main"], visit)

    assert len(batch_norm_calls) == 1
    arg_dtypes = [
        str(getattr(arg, "struct_info", getattr(arg, "ty", None)).dtype)
        for arg in batch_norm_calls[0].args
    ]
    assert arg_dtypes == ["float32"] * 5


def test_batch_norm_training_preserves_output_dtypes():
    data = helper.make_tensor_value_info("data", TensorProto.FLOAT16, [1, 3, 2, 2])
    outputs = [
        helper.make_tensor_value_info("output", TensorProto.FLOAT16, [1, 3, 2, 2]),
        helper.make_tensor_value_info("running_mean", TensorProto.FLOAT16, [3]),
        helper.make_tensor_value_info("running_var", TensorProto.FLOAT16, [3]),
    ]
    inputs = [
        data,
        helper.make_tensor_value_info("gamma", TensorProto.FLOAT16, [3]),
        helper.make_tensor_value_info("beta", TensorProto.FLOAT16, [3]),
        helper.make_tensor_value_info("mean", TensorProto.FLOAT16, [3]),
        helper.make_tensor_value_info("var", TensorProto.FLOAT16, [3]),
    ]
    batch_norm_node = helper.make_node(
        "BatchNormalization",
        [value.name for value in inputs],
        [value.name for value in outputs],
        training_mode=1,
    )
    graph = helper.make_graph(
        [batch_norm_node],
        "mixed_dtype_training_batchnorm",
        inputs,
        outputs,
    )
    model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 15)])

    tvm_model = from_onnx(model, keep_params_in_input=True)

    assert [str(field.dtype) for field in tvm_model["main"].ret_ty.fields] == [
        "float16",
        "float16",
        "float16",
    ]

    batch_norm_calls = []

    def visit(expr):
        if isinstance(expr, relax.Call) and expr.op == tvm.ir.Op.get("relax.nn.batch_norm"):
            batch_norm_calls.append(expr)

    relax.analysis.post_order_visit(tvm_model["main"], visit)

    assert len(batch_norm_calls) == 1
    assert [str(arg.ty.dtype) for arg in batch_norm_calls[0].args] == ["float32"] * 5


def get_pool_padding(shape, auto_pad, kernel_shape, strides, pads):
    def get_pad_pair(input1d, kernel1d, stride1d, mode):
        if input1d % stride1d == 0:
            pad = max(kernel1d - stride1d, 0)
        else:
            pad = max(kernel1d - (input1d % stride1d), 0)
        pad_before = pad // 2
        pad_after = pad - pad_before
        if "LOWER" in mode:
            return [pad_after, pad_before]
        return [pad_before, pad_after]

    strides = strides or [1] * (len(shape) - 2)
    padding = pads if pads is not None else 0

    if auto_pad in ("SAME_UPPER", "SAME_LOWER"):
        pad_pairs = [
            get_pad_pair(int(shape[2 + axis]), kernel_shape[axis], strides[axis], auto_pad)
            for axis in range(len(shape) - 2)
        ]
        padding = tuple(val for pair in zip(*pad_pairs) for val in pair)

    return padding


def verify_pool_ir(pool_name, shape, auto_pad, kernel_shape, strides, pads, expected):
    attrs = {
        "kernel_shape": kernel_shape,
        "strides": strides,
        "auto_pad": auto_pad,
    }
    if pads is not None:
        attrs["pads"] = pads

    node = helper.make_node(pool_name, ["x"], ["y"], **attrs)
    graph = helper.make_graph(
        [node],
        "pool_structural_test",
        inputs=[helper.make_tensor_value_info("x", TensorProto.FLOAT, shape)],
        outputs=[helper.make_tensor_value_info("y", TensorProto.FLOAT, shape)],
    )
    model = helper.make_model(graph, producer_name="pool_structural_test")
    tvm_model = from_onnx(model, keep_params_in_input=True)
    tvm.ir.assert_structural_equal(tvm_model, expected)


def test_pool():
    def make_expected(pool_name, shape, auto_pad, kernel_shape, strides, pads):
        rank = len(shape) - 2
        layout = {1: "NCW", 2: "NCHW", 3: "NCDHW"}[rank]
        padding = get_pool_padding(shape, auto_pad, kernel_shape, strides, pads)
        pool_op = {
            ("MaxPool", 1): R.nn.max_pool1d,
            ("MaxPool", 2): R.nn.max_pool2d,
            ("MaxPool", 3): R.nn.max_pool3d,
            ("AveragePool", 1): R.nn.avg_pool1d,
            ("AveragePool", 2): R.nn.avg_pool2d,
            ("AveragePool", 3): R.nn.avg_pool3d,
            ("LpPool", 1): R.nn.avg_pool1d,
            ("LpPool", 2): R.nn.avg_pool2d,
            ("LpPool", 3): R.nn.avg_pool3d,
        }[(pool_name, rank)]
        input_shape = tuple(shape)
        pool_size = kernel_shape
        dilation = [1] * rank

        if pool_name == "MaxPool":

            @I.ir_module
            class ExpectedMaxPool:
                @R.function
                def main(x: R.Tensor(input_shape, dtype="float32")):
                    R.func_attr({"num_input": 1})
                    with R.dataflow():
                        gv = pool_op(
                            x,
                            pool_size=pool_size,
                            strides=strides,
                            dilation=dilation,
                            padding=padding,
                            ceil_mode=False,
                            layout=layout,
                            out_layout=layout,
                        )
                        R.output(gv)
                    return gv

            return ExpectedMaxPool

        if pool_name == "AveragePool":

            @I.ir_module
            class ExpectedAveragePool:
                @R.function
                def main(x: R.Tensor(input_shape, dtype="float32")):
                    R.func_attr({"num_input": 1})
                    with R.dataflow():
                        gv = pool_op(
                            x,
                            pool_size=pool_size,
                            strides=strides,
                            dilation=dilation,
                            padding=padding,
                            ceil_mode=False,
                            count_include_pad=False,
                            layout=layout,
                            out_layout=layout,
                        )
                        R.output(gv)
                    return gv

            return ExpectedAveragePool

        kernel_elements = float(np.prod(kernel_shape))

        @I.ir_module
        class ExpectedLpPool:
            @R.function
            def main(x: R.Tensor(input_shape, dtype="float32")):
                R.func_attr({"num_input": 1})
                with R.dataflow():
                    lv = R.power(x, R.const(2.0, "float32"))
                    lv1 = pool_op(
                        lv,
                        pool_size=pool_size,
                        strides=strides,
                        dilation=dilation,
                        padding=padding,
                        ceil_mode=False,
                        count_include_pad=True,
                        layout=layout,
                        out_layout=layout,
                    )
                    lv2 = R.multiply(lv1, R.const(kernel_elements, "float32"))
                    gv = R.power(lv2, R.const(0.5, "float32"))
                    R.output(gv)
                return gv

        return ExpectedLpPool

    pool_cases = [
        ([1, 1, 32], "NOTSET", [3], [1], [1, 1]),
        ([1, 1, 32], "NOTSET", [3], [2], [1, 1]),
        ([1, 1, 32], "SAME_UPPER", [7], [2], None),
        ([1, 1, 32], "SAME_LOWER", [4], [4], None),
        ([1, 1, 32], "VALID", [5], [5], None),
        ([1, 1, 32], "SAME_UPPER", [3], [1], None),
        ([1, 1, 32, 32], "NOTSET", [3, 3], [1, 1], [1, 1, 1, 1]),
        ([1, 1, 32, 32], "NOTSET", [3, 3], [2, 2], [1, 1, 1, 1]),
        ([1, 1, 32, 32], "SAME_UPPER", [3, 7], [3, 2], None),
        ([1, 1, 32, 32], "SAME_LOWER", [3, 3], [2, 2], None),
        ([1, 1, 32, 32], "VALID", [3, 3], [2, 2], None),
        ([1, 1, 32, 32], "SAME_UPPER", [3, 3], [1, 1], None),
        ([1, 1, 32, 32, 32], "NOTSET", [3, 3, 4], [1, 1, 1], [1, 2, 1, 1, 2, 2]),
        ([1, 1, 32, 32, 32], "NOTSET", [3, 4, 3], [2, 2, 3], [1, 1, 1, 1, 1, 2]),
        ([1, 1, 32, 32, 32], "SAME_UPPER", [4, 3, 3], [3, 2, 2], None),
        ([1, 1, 32, 32, 32], "SAME_LOWER", [3, 3, 4], [2, 2, 2], None),
        ([1, 1, 32, 32, 32], "VALID", [3, 3, 5], [2, 2, 3], None),
        ([1, 1, 32, 32, 32], "SAME_UPPER", [3, 3, 5], [1, 1, 1], None),
    ]

    for pool_name in ["MaxPool", "AveragePool", "LpPool"]:
        for shape, auto_pad, kernel_shape, strides, pads in pool_cases:
            verify_pool_ir(
                pool_name,
                shape,
                auto_pad,
                kernel_shape,
                strides,
                pads,
                make_expected(pool_name, shape, auto_pad, kernel_shape, strides, pads),
            )


def test_global_average_pool():
    def verify_global_average_pool_ir(input_shape, expected):
        output_shape = input_shape[:2] + [1] * (len(input_shape) - 2)
        node = helper.make_node("GlobalAveragePool", ["x"], ["y"])
        graph = helper.make_graph(
            [node],
            "global_average_pool_test",
            inputs=[helper.make_tensor_value_info("x", TensorProto.FLOAT, input_shape)],
            outputs=[helper.make_tensor_value_info("y", TensorProto.FLOAT, output_shape)],
        )
        model = helper.make_model(
            graph,
            producer_name="global_average_pool_test",
            opset_imports=[helper.make_opsetid("", 14)],
        )
        tvm_model = from_onnx(model, keep_params_in_input=True)
        tvm.ir.assert_structural_equal(tvm_model, expected)

    @I.ir_module
    class Expected1D:
        @R.function
        def main(x: R.Tensor((1, 3, 32), dtype="float32")) -> R.Tensor((1, 3, 1), dtype="float32"):
            R.func_attr({"num_input": 1})
            with R.dataflow():
                gv: R.Tensor((1, 3, 1), dtype="float32") = R.mean(x, axis=[2], keepdims=True)
                R.output(gv)
            return gv

    @I.ir_module
    class Expected2D:
        @R.function
        def main(x: R.Tensor((1, 3, 32, 32), dtype="float32")) -> R.Tensor(
            (1, 3, 1, 1), dtype="float32"
        ):
            R.func_attr({"num_input": 1})
            with R.dataflow():
                gv: R.Tensor((1, 3, 1, 1), dtype="float32") = R.mean(x, axis=[2, 3], keepdims=True)
                R.output(gv)
            return gv

    @I.ir_module
    class Expected3D:
        @R.function
        def main(x: R.Tensor((1, 3, 32, 32, 32), dtype="float32")) -> R.Tensor(
            (1, 3, 1, 1, 1), dtype="float32"
        ):
            R.func_attr({"num_input": 1})
            with R.dataflow():
                gv: R.Tensor((1, 3, 1, 1, 1), dtype="float32") = R.mean(
                    x, axis=[2, 3, 4], keepdims=True
                )
                R.output(gv)
            return gv

    verify_global_average_pool_ir([1, 3, 32], Expected1D)
    verify_global_average_pool_ir([1, 3, 32, 32], Expected2D)
    verify_global_average_pool_ir([1, 3, 32, 32, 32], Expected3D)


def test_global_max_pool():
    def verify_global_max_pool_ir(input_shape, expected):
        output_shape = input_shape[:2] + [1] * (len(input_shape) - 2)
        node = helper.make_node("GlobalMaxPool", ["x"], ["y"])
        graph = helper.make_graph(
            [node],
            "global_max_pool_test",
            inputs=[helper.make_tensor_value_info("x", TensorProto.FLOAT, input_shape)],
            outputs=[helper.make_tensor_value_info("y", TensorProto.FLOAT, output_shape)],
        )
        model = helper.make_model(
            graph,
            producer_name="global_max_pool_test",
            opset_imports=[helper.make_opsetid("", 14)],
        )
        tvm_model = from_onnx(model, keep_params_in_input=True)
        tvm.ir.assert_structural_equal(tvm_model, expected)

    @I.ir_module
    class Expected1D:
        @R.function
        def main(x: R.Tensor((1, 3, 32), dtype="float32")) -> R.Tensor((1, 3, 1), dtype="float32"):
            R.func_attr({"num_input": 1})
            with R.dataflow():
                gv: R.Tensor((1, 3, 1), dtype="float32") = R.max(x, axis=[2], keepdims=True)
                R.output(gv)
            return gv

    @I.ir_module
    class Expected2D:
        @R.function
        def main(x: R.Tensor((1, 3, 32, 32), dtype="float32")) -> R.Tensor(
            (1, 3, 1, 1), dtype="float32"
        ):
            R.func_attr({"num_input": 1})
            with R.dataflow():
                gv: R.Tensor((1, 3, 1, 1), dtype="float32") = R.max(x, axis=[2, 3], keepdims=True)
                R.output(gv)
            return gv

    @I.ir_module
    class Expected3D:
        @R.function
        def main(x: R.Tensor((1, 3, 32, 32, 32), dtype="float32")) -> R.Tensor(
            (1, 3, 1, 1, 1), dtype="float32"
        ):
            R.func_attr({"num_input": 1})
            with R.dataflow():
                gv: R.Tensor((1, 3, 1, 1, 1), dtype="float32") = R.max(
                    x, axis=[2, 3, 4], keepdims=True
                )
                R.output(gv)
            return gv

    verify_global_max_pool_ir([1, 3, 32], Expected1D)
    verify_global_max_pool_ir([1, 3, 32, 32], Expected2D)
    verify_global_max_pool_ir([1, 3, 32, 32, 32], Expected3D)


@pytest.mark.parametrize("p", [1, 2, 3])
def test_global_lp_pool(p: int):
    p_value = float(p)
    inv_p_value = float(1 / p)

    def verify_global_lp_pool(input_shape, expected):
        output_shape = input_shape[:2] + [1] * (len(input_shape) - 2)
        node = helper.make_node("GlobalLpPool", ["x"], ["y"], p=p)
        graph = helper.make_graph(
            [node],
            "global_lp_pool_structural_test",
            inputs=[helper.make_tensor_value_info("x", TensorProto.FLOAT, input_shape)],
            outputs=[helper.make_tensor_value_info("y", TensorProto.FLOAT, output_shape)],
        )
        model = helper.make_model(graph, producer_name="global_lp_pool_structural_test")
        tvm_model = from_onnx(model, keep_params_in_input=True)
        tvm.ir.assert_structural_equal(tvm_model, expected)

    @I.ir_module
    class ExpectedGlobalLpPool1D:
        @R.function
        def main(
            x: R.Tensor((1, 3, 4), dtype="float32"),
        ) -> R.Tensor((1, 3, 1), dtype="float32"):
            R.func_attr({"num_input": 1})
            with R.dataflow():
                lv: R.Tensor((1, 3, 4), dtype="float32") = R.abs(x)
                lv1: R.Tensor((1, 3, 4), dtype="float32") = R.power(lv, R.const(p_value, "float32"))
                lv2: R.Tensor((1, 3, 1), dtype="float32") = R.sum(lv1, axis=[2], keepdims=True)
                gv: R.Tensor((1, 3, 1), dtype="float32") = R.power(
                    lv2, R.const(inv_p_value, "float32")
                )
                R.output(gv)
            return gv

    @I.ir_module
    class ExpectedGlobalLpPool2D:
        @R.function
        def main(
            x: R.Tensor((1, 3, 4, 4), dtype="float32"),
        ) -> R.Tensor((1, 3, 1, 1), dtype="float32"):
            R.func_attr({"num_input": 1})
            with R.dataflow():
                lv: R.Tensor((1, 3, 4, 4), dtype="float32") = R.abs(x)
                lv1: R.Tensor((1, 3, 4, 4), dtype="float32") = R.power(
                    lv, R.const(p_value, "float32")
                )
                lv2: R.Tensor((1, 3, 1, 1), dtype="float32") = R.sum(
                    lv1, axis=[2, 3], keepdims=True
                )
                gv: R.Tensor((1, 3, 1, 1), dtype="float32") = R.power(
                    lv2, R.const(inv_p_value, "float32")
                )
                R.output(gv)
            return gv

    @I.ir_module
    class ExpectedGlobalLpPool3D:
        @R.function
        def main(
            x: R.Tensor((1, 3, 4, 4, 4), dtype="float32"),
        ) -> R.Tensor((1, 3, 1, 1, 1), dtype="float32"):
            R.func_attr({"num_input": 1})
            with R.dataflow():
                lv: R.Tensor((1, 3, 4, 4, 4), dtype="float32") = R.abs(x)
                lv1: R.Tensor((1, 3, 4, 4, 4), dtype="float32") = R.power(
                    lv, R.const(p_value, "float32")
                )
                lv2: R.Tensor((1, 3, 1, 1, 1), dtype="float32") = R.sum(
                    lv1, axis=[2, 3, 4], keepdims=True
                )
                gv: R.Tensor((1, 3, 1, 1, 1), dtype="float32") = R.power(
                    lv2, R.const(inv_p_value, "float32")
                )
                R.output(gv)
            return gv

    verify_global_lp_pool([1, 3, 4], ExpectedGlobalLpPool1D)
    verify_global_lp_pool([1, 3, 4, 4], ExpectedGlobalLpPool2D)
    verify_global_lp_pool([1, 3, 4, 4, 4], ExpectedGlobalLpPool3D)


def test_maxunpool():
    input_shape = [16, 3, 16, 16]

    def verify_maxunpool(kernel_shape, pads, strides, expected):
        input_names = ["X", "I"]
        input_info = [
            helper.make_tensor_value_info("X", TensorProto.FLOAT, input_shape),
            helper.make_tensor_value_info("I", TensorProto.INT64, input_shape),
        ]

        attrs = {"kernel_shape": kernel_shape}
        if pads is not None:
            attrs["pads"] = pads
        if strides is not None:
            attrs["strides"] = strides

        node = helper.make_node("MaxUnpool", inputs=input_names, outputs=["y"], **attrs)

        graph = helper.make_graph(
            [node],
            "maxunpool_test",
            inputs=input_info,
            outputs=[helper.make_tensor_value_info("y", TensorProto.FLOAT, None)],
        )

        model = helper.make_model(graph, producer_name="maxunpool_test")
        tvm_model = from_onnx(model, keep_params_in_input=True)
        tvm.ir.assert_structural_equal(tvm_model, expected)

    @I.ir_module
    class ExpectedMaxUnpool0:
        @R.function
        def main(
            X: R.Tensor((16, 3, 16, 16), dtype="float32"),
            I_1: R.Tensor((16, 3, 16, 16), dtype="int64"),
        ) -> R.Tensor((16, 3, 17, 17), dtype="float32"):
            R.func_attr({"num_input": 2})
            with R.dataflow():
                lv: R.Tensor((16, 3, 17, 17), dtype="float32") = R.zeros(
                    R.shape([16, 3, 17, 17]), dtype="float32"
                )
                lv1: R.Tensor((13872,), dtype="float32") = R.reshape(lv, R.shape([13872]))
                lv2: R.Tensor((12288,), dtype="int64") = R.reshape(I_1, R.shape([12288]))
                lv3: R.Tensor((12288,), dtype="float32") = R.reshape(X, R.shape([12288]))
                lv4: R.Tensor((13872,), dtype="float32") = R.scatter_elements(
                    lv1, lv2, lv3, axis=0, reduction="update"
                )
                gv: R.Tensor((16, 3, 17, 17), dtype="float32") = R.reshape(
                    lv4, R.shape([16, 3, 17, 17])
                )
                R.output(gv)
            return gv

    @I.ir_module
    class ExpectedMaxUnpool1:
        @R.function
        def main(
            X: R.Tensor((16, 3, 16, 16), dtype="float32"),
            I_1: R.Tensor((16, 3, 16, 16), dtype="int64"),
        ) -> R.Tensor((16, 3, 32, 32), dtype="float32"):
            R.func_attr({"num_input": 2})
            with R.dataflow():
                lv: R.Tensor((16, 3, 32, 32), dtype="float32") = R.zeros(
                    R.shape([16, 3, 32, 32]), dtype="float32"
                )
                lv1: R.Tensor((49152,), dtype="float32") = R.reshape(lv, R.shape([49152]))
                lv2: R.Tensor((12288,), dtype="int64") = R.reshape(I_1, R.shape([12288]))
                lv3: R.Tensor((12288,), dtype="float32") = R.reshape(X, R.shape([12288]))
                lv4: R.Tensor((49152,), dtype="float32") = R.scatter_elements(
                    lv1, lv2, lv3, axis=0, reduction="update"
                )
                gv: R.Tensor((16, 3, 32, 32), dtype="float32") = R.reshape(
                    lv4, R.shape([16, 3, 32, 32])
                )
                R.output(gv)
            return gv

    @I.ir_module
    class ExpectedMaxUnpool2:
        @R.function
        def main(
            X: R.Tensor((16, 3, 16, 16), dtype="float32"),
            I_1: R.Tensor((16, 3, 16, 16), dtype="int64"),
        ) -> R.Tensor((16, 3, 15, 15), dtype="float32"):
            R.func_attr({"num_input": 2})
            with R.dataflow():
                lv: R.Tensor((16, 3, 15, 15), dtype="float32") = R.zeros(
                    R.shape([16, 3, 15, 15]), dtype="float32"
                )
                lv1: R.Tensor((10800,), dtype="float32") = R.reshape(lv, R.shape([10800]))
                lv2: R.Tensor((12288,), dtype="int64") = R.reshape(I_1, R.shape([12288]))
                lv3: R.Tensor((12288,), dtype="float32") = R.reshape(X, R.shape([12288]))
                lv4: R.Tensor((10800,), dtype="float32") = R.scatter_elements(
                    lv1, lv2, lv3, axis=0, reduction="update"
                )
                gv: R.Tensor((16, 3, 15, 15), dtype="float32") = R.reshape(
                    lv4, R.shape([16, 3, 15, 15])
                )
                R.output(gv)
            return gv

    @I.ir_module
    class ExpectedMaxUnpool3:
        @R.function
        def main(
            X: R.Tensor((16, 3, 16, 16), dtype="float32"),
            I_1: R.Tensor((16, 3, 16, 16), dtype="int64"),
        ) -> R.Tensor((16, 3, 30, 30), dtype="float32"):
            R.func_attr({"num_input": 2})
            with R.dataflow():
                lv: R.Tensor((16, 3, 30, 30), dtype="float32") = R.zeros(
                    R.shape([16, 3, 30, 30]), dtype="float32"
                )
                lv1: R.Tensor((43200,), dtype="float32") = R.reshape(lv, R.shape([43200]))
                lv2: R.Tensor((12288,), dtype="int64") = R.reshape(I_1, R.shape([12288]))
                lv3: R.Tensor((12288,), dtype="float32") = R.reshape(X, R.shape([12288]))
                lv4: R.Tensor((43200,), dtype="float32") = R.scatter_elements(
                    lv1, lv2, lv3, axis=0, reduction="update"
                )
                gv: R.Tensor((16, 3, 30, 30), dtype="float32") = R.reshape(
                    lv4, R.shape([16, 3, 30, 30])
                )
                R.output(gv)
            return gv

    @I.ir_module
    class ExpectedMaxUnpool4:
        @R.function
        def main(
            X: R.Tensor((16, 3, 16, 16), dtype="float32"),
            I_1: R.Tensor((16, 3, 16, 16), dtype="int64"),
        ) -> R.Tensor((16, 3, 18, 18), dtype="float32"):
            R.func_attr({"num_input": 2})
            with R.dataflow():
                lv: R.Tensor((16, 3, 18, 18), dtype="float32") = R.zeros(
                    R.shape([16, 3, 18, 18]), dtype="float32"
                )
                lv1: R.Tensor((15552,), dtype="float32") = R.reshape(lv, R.shape([15552]))
                lv2: R.Tensor((12288,), dtype="int64") = R.reshape(I_1, R.shape([12288]))
                lv3: R.Tensor((12288,), dtype="float32") = R.reshape(X, R.shape([12288]))
                lv4: R.Tensor((15552,), dtype="float32") = R.scatter_elements(
                    lv1, lv2, lv3, axis=0, reduction="update"
                )
                gv: R.Tensor((16, 3, 18, 18), dtype="float32") = R.reshape(
                    lv4, R.shape([16, 3, 18, 18])
                )
                R.output(gv)
            return gv

    @I.ir_module
    class ExpectedMaxUnpool5:
        @R.function
        def main(
            X: R.Tensor((16, 3, 16, 16), dtype="float32"),
            I_1: R.Tensor((16, 3, 16, 16), dtype="int64"),
        ) -> R.Tensor((16, 3, 33, 33), dtype="float32"):
            R.func_attr({"num_input": 2})
            with R.dataflow():
                lv: R.Tensor((16, 3, 33, 33), dtype="float32") = R.zeros(
                    R.shape([16, 3, 33, 33]), dtype="float32"
                )
                lv1: R.Tensor((52272,), dtype="float32") = R.reshape(lv, R.shape([52272]))
                lv2: R.Tensor((12288,), dtype="int64") = R.reshape(I_1, R.shape([12288]))
                lv3: R.Tensor((12288,), dtype="float32") = R.reshape(X, R.shape([12288]))
                lv4: R.Tensor((52272,), dtype="float32") = R.scatter_elements(
                    lv1, lv2, lv3, axis=0, reduction="update"
                )
                gv: R.Tensor((16, 3, 33, 33), dtype="float32") = R.reshape(
                    lv4, R.shape([16, 3, 33, 33])
                )
                R.output(gv)
            return gv

    @I.ir_module
    class ExpectedMaxUnpool6:
        @R.function
        def main(
            X: R.Tensor((16, 3, 16, 16), dtype="float32"),
            I_1: R.Tensor((16, 3, 16, 16), dtype="int64"),
        ) -> R.Tensor((16, 3, 16, 16), dtype="float32"):
            R.func_attr({"num_input": 2})
            with R.dataflow():
                lv: R.Tensor((16, 3, 16, 16), dtype="float32") = R.zeros(
                    R.shape([16, 3, 16, 16]), dtype="float32"
                )
                lv1: R.Tensor((12288,), dtype="float32") = R.reshape(lv, R.shape([12288]))
                lv2: R.Tensor((12288,), dtype="int64") = R.reshape(I_1, R.shape([12288]))
                lv3: R.Tensor((12288,), dtype="float32") = R.reshape(X, R.shape([12288]))
                lv4: R.Tensor((12288,), dtype="float32") = R.scatter_elements(
                    lv1, lv2, lv3, axis=0, reduction="update"
                )
                gv: R.Tensor((16, 3, 16, 16), dtype="float32") = R.reshape(
                    lv4, R.shape([16, 3, 16, 16])
                )
                R.output(gv)
            return gv

    @I.ir_module
    class ExpectedMaxUnpool7:
        @R.function
        def main(
            X: R.Tensor((16, 3, 16, 16), dtype="float32"),
            I_1: R.Tensor((16, 3, 16, 16), dtype="int64"),
        ) -> R.Tensor((16, 3, 31, 31), dtype="float32"):
            R.func_attr({"num_input": 2})
            with R.dataflow():
                lv: R.Tensor((16, 3, 31, 31), dtype="float32") = R.zeros(
                    R.shape([16, 3, 31, 31]), dtype="float32"
                )
                lv1: R.Tensor((46128,), dtype="float32") = R.reshape(lv, R.shape([46128]))
                lv2: R.Tensor((12288,), dtype="int64") = R.reshape(I_1, R.shape([12288]))
                lv3: R.Tensor((12288,), dtype="float32") = R.reshape(X, R.shape([12288]))
                lv4: R.Tensor((46128,), dtype="float32") = R.scatter_elements(
                    lv1, lv2, lv3, axis=0, reduction="update"
                )
                gv: R.Tensor((16, 3, 31, 31), dtype="float32") = R.reshape(
                    lv4, R.shape([16, 3, 31, 31])
                )
                R.output(gv)
            return gv

    verify_maxunpool([2, 2], None, None, ExpectedMaxUnpool0)
    verify_maxunpool([2, 2], None, [2, 2], ExpectedMaxUnpool1)
    verify_maxunpool([2, 2], [1, 1, 1, 1], None, ExpectedMaxUnpool2)
    verify_maxunpool([2, 2], [1, 1, 1, 1], [2, 2], ExpectedMaxUnpool3)
    verify_maxunpool([3, 3], None, None, ExpectedMaxUnpool4)
    verify_maxunpool([3, 3], None, [2, 2], ExpectedMaxUnpool5)
    verify_maxunpool([3, 3], [1, 1, 1, 1], None, ExpectedMaxUnpool6)
    verify_maxunpool([3, 3], [1, 1, 1, 1], [2, 2], ExpectedMaxUnpool7)


def test_dropout():
    def verify_dropout_ir(opset, attrs, expected):
        node = helper.make_node("Dropout", ["x"], ["y"], **attrs)
        graph = helper.make_graph(
            [node],
            "dropout_structural_test",
            inputs=[helper.make_tensor_value_info("x", TensorProto.FLOAT, [1, 3, 32, 32])],
            outputs=[helper.make_tensor_value_info("y", TensorProto.FLOAT, [1, 3, 32, 32])],
        )
        model = helper.make_model(
            graph,
            producer_name="dropout_structural_test",
            opset_imports=[helper.make_opsetid("", opset)],
        )
        tvm_model = from_onnx(model, opset=opset, keep_params_in_input=True)
        tvm.ir.assert_structural_equal(tvm_model, expected)

    @I.ir_module
    class ExpectedDropoutRateHalf:
        @R.function
        def main(x: R.Tensor((1, 3, 32, 32), dtype="float32")) -> R.Tensor(
            (1, 3, 32, 32), dtype="float32"
        ):
            R.func_attr({"num_input": 1})
            with R.dataflow():
                lv: R.Tuple(
                    R.Tensor((1, 3, 32, 32), dtype="float32"),
                    R.Tensor((1, 3, 32, 32), dtype="float32"),
                ) = R.nn.dropout(x, rate=0.5)
                lv1: R.Tensor((1, 3, 32, 32), dtype="float32") = lv[0]
                lv2: R.Tensor((1, 3, 32, 32), dtype="float32") = lv[1]
                gv: R.Tensor((1, 3, 32, 32), dtype="float32") = lv1
                R.output(gv)
            return gv

    verify_dropout_ir(14, {}, ExpectedDropoutRateHalf)
    verify_dropout_ir(11, {"ratio": 0.5}, ExpectedDropoutRateHalf)

    # Opset 12+ passes ratio as an optional input; check it is captured into the relax op.
    node = helper.make_node("Dropout", ["x", "ratio"], ["y"])
    graph = helper.make_graph(
        [node],
        "dropout_ratio_input",
        inputs=[helper.make_tensor_value_info("x", TensorProto.FLOAT, [1, 3, 4, 4])],
        initializer=[helper.make_tensor("ratio", TensorProto.FLOAT, [], [0.3])],
        outputs=[helper.make_tensor_value_info("y", TensorProto.FLOAT, [1, 3, 4, 4])],
    )
    model = helper.make_model(graph, producer_name="dropout_ratio_input")
    model.opset_import[0].version = 13
    tvm_model = from_onnx(model, opset=13, keep_params_in_input=False)

    @I.ir_module
    class ExpectedDropoutRatioInput:
        @R.function
        def main(x: R.Tensor((1, 3, 4, 4), dtype="float32")) -> R.Tensor(
            (1, 3, 4, 4), dtype="float32"
        ):
            R.func_attr({"num_input": 1})
            with R.dataflow():
                lv: R.Tuple(
                    R.Tensor((1, 3, 4, 4), dtype="float32"),
                    R.Tensor((1, 3, 4, 4), dtype="float32"),
                ) = R.nn.dropout(x, rate=0.30000001192092896)
                lv1: R.Tensor((1, 3, 4, 4), dtype="float32") = lv[0]
                lv2: R.Tensor((1, 3, 4, 4), dtype="float32") = lv[1]
                gv: R.Tensor((1, 3, 4, 4), dtype="float32") = lv1
                R.output(gv)
            return gv

    tvm.ir.assert_structural_equal(tvm_model, ExpectedDropoutRatioInput)


def test_flatten():
    def verify_flatten_ir(axis, output_shape, expected):
        node = helper.make_node("Flatten", ["x"], ["y"], axis=axis)
        graph = helper.make_graph(
            [node],
            "flatten_structural_test",
            inputs=[helper.make_tensor_value_info("x", TensorProto.FLOAT, [1, 3, 32, 32])],
            outputs=[helper.make_tensor_value_info("y", TensorProto.FLOAT, output_shape)],
        )
        model = helper.make_model(graph, producer_name="flatten_structural_test")
        tvm_model = from_onnx(model, keep_params_in_input=True)
        tvm.ir.assert_structural_equal(tvm_model, expected)

    @I.ir_module
    class ExpectedAxis0:
        @R.function
        def main(x: R.Tensor((1, 3, 32, 32), dtype="float32")) -> R.Tensor(
            (1, 3072), dtype="float32"
        ):
            R.func_attr({"num_input": 1})
            with R.dataflow():
                gv: R.Tensor((1, 3072), dtype="float32") = R.reshape(x, R.shape([1, 3072]))
                R.output(gv)
            return gv

    @I.ir_module
    class ExpectedAxisNegative1:
        @R.function
        def main(x: R.Tensor((1, 3, 32, 32), dtype="float32")) -> R.Tensor(
            (96, 32), dtype="float32"
        ):
            R.func_attr({"num_input": 1})
            with R.dataflow():
                gv: R.Tensor((96, 32), dtype="float32") = R.reshape(x, R.shape([96, 32]))
                R.output(gv)
            return gv

    @I.ir_module
    class ExpectedAxis2:
        @R.function
        def main(x: R.Tensor((1, 3, 32, 32), dtype="float32")) -> R.Tensor(
            (3, 1024), dtype="float32"
        ):
            R.func_attr({"num_input": 1})
            with R.dataflow():
                gv: R.Tensor((3, 1024), dtype="float32") = R.reshape(x, R.shape([3, 1024]))
                R.output(gv)
            return gv

    verify_flatten_ir(0, [1, 3072], ExpectedAxis0)
    verify_flatten_ir(-1, [96, 32], ExpectedAxisNegative1)
    verify_flatten_ir(2, [3, 1024], ExpectedAxis2)


def test_flatten_dynamic():
    def verify_flatten_dynamic_ir(axis, expected):
        node = helper.make_node("Flatten", ["x"], ["y"], axis=axis)
        graph = helper.make_graph(
            [node],
            "flatten_dynamic_structural_test",
            inputs=[helper.make_tensor_value_info("x", TensorProto.FLOAT, [1, "A", "B", 32])],
            outputs=[helper.make_tensor_value_info("y", TensorProto.FLOAT, [None, None])],
        )
        model = helper.make_model(graph, producer_name="flatten_dynamic_structural_test")
        tvm_model = from_onnx(model, keep_params_in_input=True)
        tvm.ir.assert_structural_equal(tvm_model, expected)

    @I.ir_module
    class ExpectedDynamicAxis0:
        @R.function
        def main(x: R.Tensor((1, "A", "B", 32), dtype="float32")) -> R.Tensor(
            (1, "A * B * 32"), dtype="float32"
        ):
            A = T.int64()
            B = T.int64()
            R.func_attr({"num_input": 1})
            with R.dataflow():
                gv: R.Tensor((1, A * B * 32), dtype="float32") = R.reshape(
                    x, R.shape([1, A * B * 32])
                )
                R.output(gv)
            return gv

    @I.ir_module
    class ExpectedDynamicAxisNegative1:
        @R.function
        def main(x: R.Tensor((1, "A", "B", 32), dtype="float32")) -> R.Tensor(
            ("A * B", 32), dtype="float32"
        ):
            A = T.int64()
            B = T.int64()
            R.func_attr({"num_input": 1})
            with R.dataflow():
                gv: R.Tensor((A * B, 32), dtype="float32") = R.reshape(x, R.shape([A * B, 32]))
                R.output(gv)
            return gv

    @I.ir_module
    class ExpectedDynamicAxis2:
        @R.function
        def main(x: R.Tensor((1, "A", "B", 32), dtype="float32")) -> R.Tensor(
            ("A", "B * 32"), dtype="float32"
        ):
            A = T.int64()
            B = T.int64()
            R.func_attr({"num_input": 1})
            with R.dataflow():
                gv: R.Tensor((A, B * 32), dtype="float32") = R.reshape(x, R.shape([A, B * 32]))
                R.output(gv)
            return gv

    verify_flatten_dynamic_ir(0, ExpectedDynamicAxis0)
    verify_flatten_dynamic_ir(-1, ExpectedDynamicAxisNegative1)
    verify_flatten_dynamic_ir(2, ExpectedDynamicAxis2)


def test_onehot():
    one_hot_node = helper.make_node("OneHot", ["indices", "depth", "values"], ["y"], axis=1)
    graph = helper.make_graph(
        [one_hot_node],
        "one_hot_test",
        inputs=[
            helper.make_tensor_value_info("indices", TensorProto.INT64, [2, 2]),
        ],
        initializer=[
            helper.make_tensor("depth", TensorProto.INT64, [], [10]),
            helper.make_tensor("values", TensorProto.FLOAT, [2], [3, 1]),
        ],
        outputs=[helper.make_tensor_value_info("y", TensorProto.FLOAT, [2, 10, 2])],
    )

    model = helper.make_model(graph, producer_name="one_hot_test")
    values = {
        "indices": np.array([[1, 9], [2, 4]], dtype="int64"),
    }
    check_correctness(model, inputs=values)


@pytest.mark.parametrize("axis", [None, 0, 1, -1])
@pytest.mark.parametrize("sorted", [0, 1])
@pytest.mark.parametrize("num_outputs", [1, 2, 3, 4])
def test_unique(axis: int | None, sorted: int, num_outputs: int):
    if num_outputs in [3, 4] and axis is None:
        pytest.xfail("RuntimeError: Check failed: input_shape.size() == size (2 vs. 1)")

    input_shape = [8, 8]
    if axis is None:
        output_shape = [-1]
    else:
        output_shape = [8, 8]
        output_shape[axis] = -1

    output_names = ["y", "indices", "inverse_indices", "counts"][:num_outputs]
    unique_node = helper.make_node("Unique", ["x"], output_names, axis=axis, sorted=sorted)

    outputs = [helper.make_tensor_value_info("y", TensorProto.FLOAT, output_shape)]
    if num_outputs > 1:
        outputs.append(helper.make_tensor_value_info("indices", TensorProto.INT64, [-1]))
    if num_outputs > 2:
        # ONNX spec: inverse_indices is always 1D
        outputs.append(helper.make_tensor_value_info("inverse_indices", TensorProto.INT64, [-1]))
    if num_outputs > 3:
        outputs.append(helper.make_tensor_value_info("counts", TensorProto.INT64, [-1]))

    graph = helper.make_graph(
        [unique_node],
        "unique_test",
        inputs=[helper.make_tensor_value_info("x", TensorProto.FLOAT, input_shape)],
        outputs=outputs,
    )
    model = helper.make_model(graph, producer_name="unique_test")
    check_correctness(model)


def test_nonzero():
    def verify_nonzero(shape, expected):
        ndim = max(len(shape), 1)
        node = helper.make_node("NonZero", ["x"], ["y"])
        graph = helper.make_graph(
            [node],
            "nonzero_structural_test",
            inputs=[helper.make_tensor_value_info("x", TensorProto.BOOL, shape)],
            outputs=[helper.make_tensor_value_info("y", TensorProto.INT64, [ndim, None])],
        )
        model = helper.make_model(graph, producer_name="nonzero_structural_test")
        tvm_model = from_onnx(model, keep_params_in_input=True)
        tvm.ir.assert_structural_equal(tvm_model, expected)

    @I.ir_module
    class ExpectedScalar:
        @R.function
        def main(x: R.Tensor((), dtype="bool")):
            nonzero_numbers = T.int64()
            R.func_attr({"num_input": 1})
            with R.dataflow():
                lv: R.Tensor((1, nonzero_numbers), dtype="int64") = R.match_cast(
                    R.nonzero(x), R.Tensor((1, nonzero_numbers), dtype="int64")
                )
                gv: R.Tensor((1, nonzero_numbers), dtype="int64") = lv
                R.output(gv)
            return gv

    @I.ir_module
    class ExpectedRank1:
        @R.function
        def main(x: R.Tensor((1,), dtype="bool")):
            nonzero_numbers = T.int64()
            R.func_attr({"num_input": 1})
            with R.dataflow():
                lv: R.Tensor((1, nonzero_numbers), dtype="int64") = R.match_cast(
                    R.nonzero(x), R.Tensor((1, nonzero_numbers), dtype="int64")
                )
                gv: R.Tensor((1, nonzero_numbers), dtype="int64") = lv
                R.output(gv)
            return gv

    @I.ir_module
    class ExpectedRank2:
        @R.function
        def main(x: R.Tensor((2, 3), dtype="bool")):
            nonzero_numbers = T.int64()
            R.func_attr({"num_input": 1})
            with R.dataflow():
                lv: R.Tensor((2, nonzero_numbers), dtype="int64") = R.match_cast(
                    R.nonzero(x), R.Tensor((2, nonzero_numbers), dtype="int64")
                )
                gv: R.Tensor((2, nonzero_numbers), dtype="int64") = lv
                R.output(gv)
            return gv

    @I.ir_module
    class ExpectedRank3:
        @R.function
        def main(x: R.Tensor((4, 5, 6), dtype="bool")):
            nonzero_numbers = T.int64()
            R.func_attr({"num_input": 1})
            with R.dataflow():
                lv: R.Tensor((3, nonzero_numbers), dtype="int64") = R.match_cast(
                    R.nonzero(x), R.Tensor((3, nonzero_numbers), dtype="int64")
                )
                gv: R.Tensor((3, nonzero_numbers), dtype="int64") = lv
                R.output(gv)
            return gv

    @I.ir_module
    class ExpectedRank4:
        @R.function
        def main(x: R.Tensor((7, 8, 9, 10), dtype="bool")):
            nonzero_numbers = T.int64()
            R.func_attr({"num_input": 1})
            with R.dataflow():
                lv: R.Tensor((4, nonzero_numbers), dtype="int64") = R.match_cast(
                    R.nonzero(x), R.Tensor((4, nonzero_numbers), dtype="int64")
                )
                gv: R.Tensor((4, nonzero_numbers), dtype="int64") = lv
                R.output(gv)
            return gv

    verify_nonzero((), ExpectedScalar)
    verify_nonzero((1,), ExpectedRank1)
    verify_nonzero((2, 3), ExpectedRank2)
    verify_nonzero((4, 5, 6), ExpectedRank3)
    verify_nonzero((7, 8, 9, 10), ExpectedRank4)


def test_depth_to_space():
    def verify_depth_to_space(mode: Literal["DCR", "CRD"], expected):
        in_shape = [1, 8, 2, 3]
        out_shape = [1, 2, 4, 6]
        blocksize = 2
        node = onnx.helper.make_node(
            "DepthToSpace", inputs=["x"], outputs=["y"], blocksize=blocksize, mode=mode
        )
        graph = helper.make_graph(
            [node],
            "depth_to_space_test",
            inputs=[helper.make_tensor_value_info("x", TensorProto.FLOAT, in_shape)],
            outputs=[helper.make_tensor_value_info("y", TensorProto.FLOAT, out_shape)],
        )
        model = helper.make_model(graph, producer_name="depth_to_space_test")
        tvm_model = from_onnx(model, keep_params_in_input=True)
        tvm.ir.assert_structural_equal(tvm_model, expected)

    @I.ir_module
    class ExpectedDCR:
        @R.function
        def main(
            x: R.Tensor((1, 8, 2, 3), dtype="float32"),
        ) -> R.Tensor((1, 2, 4, 6), dtype="float32"):
            R.func_attr({"num_input": 1})
            with R.dataflow():
                lv: R.Tensor((1, 2, 2, 2, 2, 3), dtype="float32") = R.reshape(
                    x, R.shape([1, 2, 2, 2, 2, 3])
                )
                lv1: R.Tensor((1, 2, 2, 2, 3, 2), dtype="float32") = R.permute_dims(
                    lv, axes=[0, 3, 4, 1, 5, 2]
                )
                gv: R.Tensor((1, 2, 4, 6), dtype="float32") = R.reshape(lv1, R.shape([1, 2, 4, 6]))
                R.output(gv)
            return gv

    @I.ir_module
    class ExpectedCRD:
        @R.function
        def main(
            x: R.Tensor((1, 8, 2, 3), dtype="float32"),
        ) -> R.Tensor((1, 2, 4, 6), dtype="float32"):
            R.func_attr({"num_input": 1})
            with R.dataflow():
                lv: R.Tensor((1, 2, 2, 2, 2, 3), dtype="float32") = R.reshape(
                    x, R.shape([1, 2, 2, 2, 2, 3])
                )
                lv1: R.Tensor((1, 2, 2, 2, 3, 2), dtype="float32") = R.permute_dims(
                    lv, axes=[0, 1, 4, 2, 5, 3]
                )
                gv: R.Tensor((1, 2, 4, 6), dtype="float32") = R.reshape(lv1, R.shape([1, 2, 4, 6]))
                R.output(gv)
            return gv

    verify_depth_to_space("DCR", ExpectedDCR)
    verify_depth_to_space("CRD", ExpectedCRD)


def test_space_to_depth():
    in_shape = [1, 2, 4, 6]
    out_shape = [1, 8, 2, 3]
    blocksize = 2
    node = onnx.helper.make_node("SpaceToDepth", inputs=["x"], outputs=["y"], blocksize=blocksize)
    graph = helper.make_graph(
        [node],
        "space_to_depth_test",
        inputs=[helper.make_tensor_value_info("x", TensorProto.FLOAT, in_shape)],
        outputs=[helper.make_tensor_value_info("y", TensorProto.FLOAT, out_shape)],
    )
    model = helper.make_model(graph, producer_name="space_to_depth_test")
    tvm_model = from_onnx(model, keep_params_in_input=True)

    @I.ir_module
    class Expected:
        @R.function
        def main(
            x: R.Tensor((1, 2, 4, 6), dtype="float32"),
        ) -> R.Tensor((1, 8, 2, 3), dtype="float32"):
            R.func_attr({"num_input": 1})
            with R.dataflow():
                lv: R.Tensor((1, 2, 2, 2, 3, 2), dtype="float32") = R.reshape(
                    x, R.shape([1, 2, 2, 2, 3, 2])
                )
                lv1: R.Tensor((1, 2, 2, 2, 2, 3), dtype="float32") = R.permute_dims(
                    lv, axes=[0, 3, 5, 1, 2, 4]
                )
                gv: R.Tensor((1, 8, 2, 3), dtype="float32") = R.reshape(lv1, R.shape([1, 8, 2, 3]))
                R.output(gv)
            return gv

    tvm.ir.assert_structural_equal(tvm_model, Expected)


def construct_sequence(input_shape: list[int], num_tensors: int, name: str = "sequence"):
    inputs = [f"data{i}" for i in range(num_tensors)]
    sequence_construct_node = helper.make_node("SequenceConstruct", inputs, [name])
    graph_inputs = [
        helper.make_tensor_value_info(f"data{i}", TensorProto.FLOAT, input_shape)
        for i in range(num_tensors)
    ]
    return sequence_construct_node, graph_inputs


def make_constant_node(name: str, data_type: int, dims: list[int], vals: list[int]):
    return helper.make_node(
        "Constant",
        inputs=[],
        outputs=[name],
        value=helper.make_tensor(name=name, data_type=data_type, dims=dims, vals=vals),
    )


def make_optional_tensor_value_info(name: str, elem_type: int, shape: list[int]):
    return helper.make_value_info(
        name, helper.make_optional_type_proto(helper.make_tensor_type_proto(elem_type, shape))
    )


def make_optional_sequence_value_info(name: str, elem_type: int, shape: list[int]):
    return helper.make_value_info(
        name,
        helper.make_optional_type_proto(
            helper.make_sequence_type_proto(helper.make_tensor_type_proto(elem_type, shape))
        ),
    )


def test_sequence_construct():
    node, graph_inputs = construct_sequence(input_shape=[32, 32], num_tensors=2)
    graph = helper.make_graph(
        [node],
        "test_sequence_construct",
        inputs=graph_inputs,
        outputs=[helper.make_tensor_sequence_value_info("sequence", TensorProto.FLOAT, [32, 32])],
    )
    model = helper.make_model(graph, producer_name="test_sequence_construct")
    tvm_model = from_onnx(model, keep_params_in_input=True)

    @I.ir_module
    class Expected:
        @R.function
        def main(
            data0: R.Tensor((32, 32), dtype="float32"),
            data1: R.Tensor((32, 32), dtype="float32"),
        ) -> R.Tuple(R.Tensor((32, 32), dtype="float32"), R.Tensor((32, 32), dtype="float32")):
            R.func_attr({"num_input": 2})
            with R.dataflow():
                gv: R.Tuple(
                    R.Tensor((32, 32), dtype="float32"),
                    R.Tensor((32, 32), dtype="float32"),
                ) = data0, data1
                R.output(gv)
            return gv

    tvm.ir.assert_structural_equal(tvm_model, Expected)


def test_sequence_empty():
    sequence_empty_node = helper.make_node("SequenceEmpty", [], ["sequence"])
    graph = helper.make_graph(
        [sequence_empty_node],
        "test_sequence_empty",
        inputs=[],
        outputs=[helper.make_tensor_sequence_value_info("sequence", TensorProto.FLOAT, [])],
    )
    model = helper.make_model(graph, producer_name="test_sequence_empty")
    tvm_model = from_onnx(model, keep_params_in_input=True)

    @I.ir_module
    class Expected:
        @R.function
        def main() -> R.Tuple:
            R.func_attr({"num_input": 0})
            with R.dataflow():
                gv: R.Tuple = R.tuple()
                R.output(gv)
            return R.tuple()

    tvm.ir.assert_structural_equal(tvm_model, Expected)


def test_sequence_erase():
    def verify_sequence_erase(explicit_position: bool, expected):
        seq_node, graph_inputs = construct_sequence(input_shape=[32, 32], num_tensors=4)
        index = make_constant_node("index", TensorProto.INT64, (), [1])
        node_input = ["sequence", "index"] if explicit_position else ["sequence"]
        sequence_erase_node = helper.make_node("SequenceErase", node_input, ["output"])
        graph = helper.make_graph(
            [index, seq_node, sequence_erase_node],
            "test_sequence_erase",
            inputs=graph_inputs,
            outputs=[helper.make_tensor_sequence_value_info("output", TensorProto.FLOAT, [32, 32])],
        )
        model = helper.make_model(graph, producer_name="test_sequence_erase")
        tvm_model = from_onnx(model, keep_params_in_input=True)
        tvm.ir.assert_structural_equal(tvm_model, expected)

    @I.ir_module
    class ExpectedEraseExplicit:
        @R.function
        def main(
            data0: R.Tensor((32, 32), dtype="float32"),
            data1: R.Tensor((32, 32), dtype="float32"),
            data2: R.Tensor((32, 32), dtype="float32"),
            data3: R.Tensor((32, 32), dtype="float32"),
        ) -> R.Tuple(
            R.Tensor((32, 32), dtype="float32"),
            R.Tensor((32, 32), dtype="float32"),
            R.Tensor((32, 32), dtype="float32"),
        ):
            R.func_attr({"num_input": 4})
            with R.dataflow():
                gv: R.Tuple(
                    R.Tensor((32, 32), dtype="float32"),
                    R.Tensor((32, 32), dtype="float32"),
                    R.Tensor((32, 32), dtype="float32"),
                ) = data0, data2, data3
                R.output(gv)
            return gv

    @I.ir_module
    class ExpectedEraseDefault:
        @R.function
        def main(
            data0: R.Tensor((32, 32), dtype="float32"),
            data1: R.Tensor((32, 32), dtype="float32"),
            data2: R.Tensor((32, 32), dtype="float32"),
            data3: R.Tensor((32, 32), dtype="float32"),
        ) -> R.Tuple(
            R.Tensor((32, 32), dtype="float32"),
            R.Tensor((32, 32), dtype="float32"),
            R.Tensor((32, 32), dtype="float32"),
        ):
            R.func_attr({"num_input": 4})
            with R.dataflow():
                gv: R.Tuple(
                    R.Tensor((32, 32), dtype="float32"),
                    R.Tensor((32, 32), dtype="float32"),
                    R.Tensor((32, 32), dtype="float32"),
                ) = data0, data1, data2
                R.output(gv)
            return gv

    verify_sequence_erase(True, ExpectedEraseExplicit)
    verify_sequence_erase(False, ExpectedEraseDefault)


def test_sequence_insert():
    def verify_sequence_insert(explicit_position: bool, expected):
        seq_node, graph_inputs = construct_sequence(input_shape=[32, 32], num_tensors=4)
        index = make_constant_node("index", TensorProto.INT64, (), [0])
        node_input = ["sequence", "value", "index"] if explicit_position else ["sequence", "value"]
        sequence_insert_node = helper.make_node("SequenceInsert", node_input, ["output"])
        graph = helper.make_graph(
            [index, seq_node, sequence_insert_node],
            "test_sequence_insert",
            inputs=[
                *graph_inputs,
                helper.make_tensor_value_info("value", TensorProto.FLOAT, [32, 32]),
            ],
            outputs=[helper.make_tensor_sequence_value_info("output", TensorProto.FLOAT, [32, 32])],
        )
        model = helper.make_model(graph, producer_name="test_sequence_insert")
        tvm_model = from_onnx(model, keep_params_in_input=True)
        tvm.ir.assert_structural_equal(tvm_model, expected)

    @I.ir_module
    class ExpectedInsertExplicit:
        @R.function
        def main(
            data0: R.Tensor((32, 32), dtype="float32"),
            data1: R.Tensor((32, 32), dtype="float32"),
            data2: R.Tensor((32, 32), dtype="float32"),
            data3: R.Tensor((32, 32), dtype="float32"),
            value: R.Tensor((32, 32), dtype="float32"),
        ) -> R.Tuple(
            R.Tensor((32, 32), dtype="float32"),
            R.Tensor((32, 32), dtype="float32"),
            R.Tensor((32, 32), dtype="float32"),
            R.Tensor((32, 32), dtype="float32"),
            R.Tensor((32, 32), dtype="float32"),
        ):
            R.func_attr({"num_input": 5})
            with R.dataflow():
                gv: R.Tuple(
                    R.Tensor((32, 32), dtype="float32"),
                    R.Tensor((32, 32), dtype="float32"),
                    R.Tensor((32, 32), dtype="float32"),
                    R.Tensor((32, 32), dtype="float32"),
                    R.Tensor((32, 32), dtype="float32"),
                ) = value, data0, data1, data2, data3
                R.output(gv)
            return gv

    @I.ir_module
    class ExpectedInsertDefault:
        @R.function
        def main(
            data0: R.Tensor((32, 32), dtype="float32"),
            data1: R.Tensor((32, 32), dtype="float32"),
            data2: R.Tensor((32, 32), dtype="float32"),
            data3: R.Tensor((32, 32), dtype="float32"),
            value: R.Tensor((32, 32), dtype="float32"),
        ) -> R.Tuple(
            R.Tensor((32, 32), dtype="float32"),
            R.Tensor((32, 32), dtype="float32"),
            R.Tensor((32, 32), dtype="float32"),
            R.Tensor((32, 32), dtype="float32"),
            R.Tensor((32, 32), dtype="float32"),
        ):
            R.func_attr({"num_input": 5})
            with R.dataflow():
                gv: R.Tuple(
                    R.Tensor((32, 32), dtype="float32"),
                    R.Tensor((32, 32), dtype="float32"),
                    R.Tensor((32, 32), dtype="float32"),
                    R.Tensor((32, 32), dtype="float32"),
                    R.Tensor((32, 32), dtype="float32"),
                ) = data0, data1, data2, data3, value
                R.output(gv)
            return gv

    verify_sequence_insert(True, ExpectedInsertExplicit)
    verify_sequence_insert(False, ExpectedInsertDefault)


def test_concat_from_sequence():
    def verify_concat_from_sequence(new_axis: int, axis: int, expected_shape: list[int], expected):
        seq_node, graph_inputs = construct_sequence(input_shape=[32, 32], num_tensors=2)
        concat_from_sequence_node = helper.make_node(
            "ConcatFromSequence", ["sequence"], ["output"], axis=axis, new_axis=new_axis
        )
        graph = helper.make_graph(
            [seq_node, concat_from_sequence_node],
            "test_concat_from_sequence",
            inputs=graph_inputs,
            outputs=[helper.make_tensor_value_info("output", TensorProto.FLOAT, expected_shape)],
        )
        model = helper.make_model(graph, producer_name="test_concat_from_sequence")
        tvm_model = from_onnx(model, keep_params_in_input=True)
        tvm.ir.assert_structural_equal(tvm_model, expected)

    @I.ir_module
    class ExpectedConcatAxis0:
        @R.function
        def main(
            data0: R.Tensor((32, 32), dtype="float32"),
            data1: R.Tensor((32, 32), dtype="float32"),
        ) -> R.Tensor((64, 32), dtype="float32"):
            R.func_attr({"num_input": 2})
            with R.dataflow():
                gv: R.Tensor((64, 32), dtype="float32") = R.concat((data0, data1), axis=0)
                R.output(gv)
            return gv

    @I.ir_module
    class ExpectedConcatAxis1:
        @R.function
        def main(
            data0: R.Tensor((32, 32), dtype="float32"),
            data1: R.Tensor((32, 32), dtype="float32"),
        ) -> R.Tensor((32, 64), dtype="float32"):
            R.func_attr({"num_input": 2})
            with R.dataflow():
                gv: R.Tensor((32, 64), dtype="float32") = R.concat((data0, data1), axis=1)
                R.output(gv)
            return gv

    @I.ir_module
    class ExpectedStackAxis0:
        @R.function
        def main(
            data0: R.Tensor((32, 32), dtype="float32"),
            data1: R.Tensor((32, 32), dtype="float32"),
        ) -> R.Tensor((2, 32, 32), dtype="float32"):
            R.func_attr({"num_input": 2})
            with R.dataflow():
                lv: R.Tensor((1, 32, 32), dtype="float32") = R.expand_dims(data0, axis=[0])
                lv1: R.Tensor((1, 32, 32), dtype="float32") = R.expand_dims(data1, axis=[0])
                gv: R.Tensor((2, 32, 32), dtype="float32") = R.concat((lv, lv1), axis=0)
                R.output(gv)
            return gv

    @I.ir_module
    class ExpectedStackAxis1:
        @R.function
        def main(
            data0: R.Tensor((32, 32), dtype="float32"),
            data1: R.Tensor((32, 32), dtype="float32"),
        ) -> R.Tensor((32, 2, 32), dtype="float32"):
            R.func_attr({"num_input": 2})
            with R.dataflow():
                lv: R.Tensor((32, 1, 32), dtype="float32") = R.expand_dims(data0, axis=[1])
                lv1: R.Tensor((32, 1, 32), dtype="float32") = R.expand_dims(data1, axis=[1])
                gv: R.Tensor((32, 2, 32), dtype="float32") = R.concat((lv, lv1), axis=1)
                R.output(gv)
            return gv

    @I.ir_module
    class ExpectedStackAxisMinusOne:
        @R.function
        def main(
            data0: R.Tensor((32, 32), dtype="float32"),
            data1: R.Tensor((32, 32), dtype="float32"),
        ) -> R.Tensor((32, 32, 2), dtype="float32"):
            R.func_attr({"num_input": 2})
            with R.dataflow():
                lv: R.Tensor((32, 32, 1), dtype="float32") = R.expand_dims(data0, axis=[-1])
                lv1: R.Tensor((32, 32, 1), dtype="float32") = R.expand_dims(data1, axis=[-1])
                gv: R.Tensor((32, 32, 2), dtype="float32") = R.concat((lv, lv1), axis=-1)
                R.output(gv)
            return gv

    verify_concat_from_sequence(0, 0, [64, 32], ExpectedConcatAxis0)
    verify_concat_from_sequence(0, 1, [32, 64], ExpectedConcatAxis1)
    verify_concat_from_sequence(1, 0, [2, 32, 32], ExpectedStackAxis0)
    verify_concat_from_sequence(1, 1, [32, 2, 32], ExpectedStackAxis1)
    verify_concat_from_sequence(1, -1, [32, 32, 2], ExpectedStackAxisMinusOne)


def test_concat_from_sequence_new_axis_three_tensors():
    """new_axis=1 with three sequence elements (stack then concat along axis)."""
    seq_node, graph_inputs = construct_sequence(input_shape=[16, 8], num_tensors=3)
    concat_node = helper.make_node(
        "ConcatFromSequence", ["sequence"], ["output"], axis=0, new_axis=1
    )
    graph = helper.make_graph(
        [seq_node, concat_node],
        "test_concat_from_sequence_new_axis_three",
        inputs=graph_inputs,
        outputs=[helper.make_tensor_value_info("output", TensorProto.FLOAT, [3, 16, 8])],
    )
    model = helper.make_model(graph, producer_name="test_concat_from_sequence_new_axis_three")
    tvm_model = from_onnx(model, keep_params_in_input=True)

    @I.ir_module
    class Expected:
        @R.function
        def main(
            data0: R.Tensor((16, 8), dtype="float32"),
            data1: R.Tensor((16, 8), dtype="float32"),
            data2: R.Tensor((16, 8), dtype="float32"),
        ) -> R.Tensor((3, 16, 8), dtype="float32"):
            R.func_attr({"num_input": 3})
            with R.dataflow():
                lv: R.Tensor((1, 16, 8), dtype="float32") = R.expand_dims(data0, axis=[0])
                lv1: R.Tensor((1, 16, 8), dtype="float32") = R.expand_dims(data1, axis=[0])
                lv2: R.Tensor((1, 16, 8), dtype="float32") = R.expand_dims(data2, axis=[0])
                gv: R.Tensor((3, 16, 8), dtype="float32") = R.concat((lv, lv1, lv2), axis=0)
                R.output(gv)
            return gv

    tvm.ir.assert_structural_equal(tvm_model, Expected)


def test_concat_from_sequence_invalid_new_axis():
    """Verify that new_axis values other than 0 or 1 raise a ValueError."""
    seq_node, graph_inputs = construct_sequence(input_shape=[16, 8], num_tensors=2)
    concat_node = helper.make_node(
        "ConcatFromSequence", ["sequence"], ["output"], axis=0, new_axis=2
    )
    graph = helper.make_graph(
        [seq_node, concat_node],
        "test_concat_from_sequence_invalid_new_axis",
        inputs=graph_inputs,
        outputs=[helper.make_tensor_value_info("output", TensorProto.FLOAT, [32, 8])],
    )
    model = helper.make_model(graph, producer_name="test_concat_from_sequence_invalid_new_axis")

    with pytest.raises(ValueError, match="ConcatFromSequence only supports new_axis in"):
        from_onnx(model, opset=11)


def test_split_to_sequence():
    def verify_split_to_sequence(split, data_shape: list[int], output_shape: list[int], expected):
        split_to_sequence_node = helper.make_node(
            "SplitToSequence",
            ["data", "split"],
            ["output"],
            axis=0,
        )
        split_shape = [len(split)] if isinstance(split, list) else ()
        split_node = make_constant_node(
            "split", TensorProto.INT64, split_shape, [split] if isinstance(split, int) else split
        )
        graph = helper.make_graph(
            [split_node, split_to_sequence_node],
            "test_split_to_sequence",
            inputs=[helper.make_tensor_value_info("data", TensorProto.FLOAT, data_shape)],
            outputs=[
                helper.make_tensor_sequence_value_info("output", TensorProto.FLOAT, output_shape)
            ],
        )
        model = helper.make_model(graph, producer_name="test_split_to_sequence")
        tvm_model = from_onnx(model, keep_params_in_input=True)
        tvm.ir.assert_structural_equal(tvm_model, expected)

    @I.ir_module
    class ExpectedScalarSplit:
        @R.function
        def main(
            data: R.Tensor((6, 32), dtype="float32"),
        ) -> R.Tuple(
            R.Tensor((2, 32), dtype="float32"),
            R.Tensor((2, 32), dtype="float32"),
            R.Tensor((2, 32), dtype="float32"),
        ):
            R.func_attr({"num_input": 1})
            with R.dataflow():
                gv: R.Tuple(
                    R.Tensor((2, 32), dtype="float32"),
                    R.Tensor((2, 32), dtype="float32"),
                    R.Tensor((2, 32), dtype="float32"),
                ) = R.split(data, indices_or_sections=3, axis=0)
                R.output(gv)
            return gv

    @I.ir_module
    class ExpectedSectionsSplit:
        @R.function
        def main(
            data: R.Tensor((64, 32), dtype="float32"),
        ) -> R.Tuple(
            R.Tensor((16, 32), dtype="float32"),
            R.Tensor((48, 32), dtype="float32"),
        ):
            R.func_attr({"num_input": 1})
            with R.dataflow():
                gv: R.Tuple(
                    R.Tensor((16, 32), dtype="float32"),
                    R.Tensor((48, 32), dtype="float32"),
                ) = R.split(data, indices_or_sections=[16], axis=0)
                R.output(gv)
            return gv

    verify_split_to_sequence(2, [6, 32], [2, 32], ExpectedScalarSplit)
    verify_split_to_sequence([16, 48], [64, 32], [32, 32], ExpectedSectionsSplit)


def test_sequence_at():
    seq_node, graph_inputs = construct_sequence(input_shape=[32, 32], num_tensors=4)
    index = make_constant_node("index", TensorProto.INT64, (), [1])
    node_input = ["sequence", "index"]
    sequence_at_node = helper.make_node("SequenceAt", node_input, ["output"])
    graph = helper.make_graph(
        [index, seq_node, sequence_at_node],
        "test_sequence_at",
        inputs=graph_inputs,
        outputs=[helper.make_tensor_value_info("output", TensorProto.FLOAT, [32, 32])],
    )
    model = helper.make_model(graph, producer_name="test_sequence_at")
    tvm_model = from_onnx(model, keep_params_in_input=True)

    @I.ir_module
    class Expected:
        @R.function
        def main(
            data0: R.Tensor((32, 32), dtype="float32"),
            data1: R.Tensor((32, 32), dtype="float32"),
            data2: R.Tensor((32, 32), dtype="float32"),
            data3: R.Tensor((32, 32), dtype="float32"),
        ) -> R.Tensor((32, 32), dtype="float32"):
            R.func_attr({"num_input": 4})
            with R.dataflow():
                gv: R.Tensor((32, 32), dtype="float32") = data1
                R.output(gv)
            return gv

    tvm.ir.assert_structural_equal(tvm_model, Expected)


def test_optional_get_element_tensor():
    x_shape = [2, 3]
    optional_node = helper.make_node("Optional", ["x"], ["optional"])
    get_element_node = helper.make_node("OptionalGetElement", ["optional"], ["output"])
    graph = helper.make_graph(
        [optional_node, get_element_node],
        "test_optional_get_element_tensor",
        inputs=[helper.make_tensor_value_info("x", TensorProto.FLOAT, x_shape)],
        outputs=[helper.make_tensor_value_info("output", TensorProto.FLOAT, x_shape)],
        value_info=[make_optional_tensor_value_info("optional", TensorProto.FLOAT, x_shape)],
    )
    model = helper.make_model(graph, producer_name="test_optional_get_element_tensor")
    model.ir_version = 11
    model.opset_import[0].version = 18
    tvm_model = from_onnx(model, opset=18, keep_params_in_input=True)

    @I.ir_module
    class Expected:
        @R.function
        def main(x: R.Tensor((2, 3), dtype="float32")) -> R.Tensor((2, 3), dtype="float32"):
            R.func_attr({"num_input": 1})
            with R.dataflow():
                gv: R.Tensor((2, 3), dtype="float32") = x
                R.output(gv)
            return gv

    tvm.ir.assert_structural_equal(tvm_model, Expected)


def test_optional_has_element_tensor():
    x_shape = [2, 3]
    optional_node = helper.make_node("Optional", ["x"], ["optional"])
    has_element_node = helper.make_node("OptionalHasElement", ["optional"], ["output"])
    graph = helper.make_graph(
        [optional_node, has_element_node],
        "test_optional_has_element_tensor",
        inputs=[helper.make_tensor_value_info("x", TensorProto.FLOAT, x_shape)],
        outputs=[helper.make_tensor_value_info("output", TensorProto.BOOL, [])],
        value_info=[make_optional_tensor_value_info("optional", TensorProto.FLOAT, x_shape)],
    )
    model = helper.make_model(graph, producer_name="test_optional_has_element_tensor")
    model.ir_version = 11
    model.opset_import[0].version = 18
    tvm_model = from_onnx(model, opset=18, keep_params_in_input=True)

    @I.ir_module
    class Expected:
        @R.function
        def main(x: R.Tensor((2, 3), dtype="float32")) -> R.Tensor((), dtype="bool"):
            R.func_attr({"num_input": 1})
            with R.dataflow():
                gv: R.Tensor((), dtype="bool") = R.const(True, "bool")
                R.output(gv)
            return gv

    tvm.ir.assert_structural_equal(tvm_model, Expected)


def test_optional_has_element_empty():
    x_shape = [2, 3]
    tensor_type = helper.make_tensor_type_proto(TensorProto.FLOAT, x_shape)
    optional_type = helper.make_optional_type_proto(tensor_type)
    optional_node = helper.make_node("Optional", [], ["optional"], type=tensor_type)
    has_element_node = helper.make_node("OptionalHasElement", ["optional"], ["output"])
    graph = helper.make_graph(
        [optional_node, has_element_node],
        "test_optional_has_element_empty",
        inputs=[],
        outputs=[helper.make_tensor_value_info("output", TensorProto.BOOL, [])],
        value_info=[helper.make_value_info("optional", optional_type)],
    )
    model = helper.make_model(graph, producer_name="test_optional_has_element_empty")
    model.ir_version = 11
    model.opset_import[0].version = 18
    tvm_model = from_onnx(model, opset=18, keep_params_in_input=True)

    @I.ir_module
    class Expected:
        @R.function
        def main() -> R.Tensor((), dtype="bool"):
            R.func_attr({"num_input": 0})
            with R.dataflow():
                gv: R.Tensor((), dtype="bool") = R.const(False, "bool")
                R.output(gv)
            return gv

    tvm.ir.assert_structural_equal(tvm_model, Expected)


def test_optional_has_element_empty_ir():
    x_shape = [2, 3]
    tensor_type = helper.make_tensor_type_proto(TensorProto.FLOAT, x_shape)
    optional_type = helper.make_optional_type_proto(tensor_type)
    optional_node = helper.make_node("Optional", [], ["optional"], type=tensor_type)
    has_element_node = helper.make_node("OptionalHasElement", ["optional"], ["output"])
    graph = helper.make_graph(
        [optional_node, has_element_node],
        "test_optional_has_element_empty_ir",
        inputs=[],
        outputs=[helper.make_tensor_value_info("output", TensorProto.BOOL, [])],
        value_info=[helper.make_value_info("optional", optional_type)],
    )
    model = helper.make_model(graph, producer_name="test_optional_has_element_empty_ir")
    model.ir_version = 11
    model.opset_import[0].version = 18
    tvm_model = from_onnx(model, opset=18, keep_params_in_input=True)

    @I.ir_module
    class Expected:
        @R.function
        def main() -> R.Tensor((), dtype="bool"):
            R.func_attr({"num_input": 0})
            with R.dataflow():
                gv: R.Tensor((), dtype="bool") = R.const(False, "bool")
                R.output(gv)
            return gv

    tvm.ir.assert_structural_equal(tvm_model, Expected)


def test_optional_get_element_tensor_ir():
    x_shape = [2, 3]
    optional_node = helper.make_node("Optional", ["x"], ["optional"])
    get_element_node = helper.make_node("OptionalGetElement", ["optional"], ["output"])
    graph = helper.make_graph(
        [optional_node, get_element_node],
        "test_optional_get_element_tensor_ir",
        inputs=[helper.make_tensor_value_info("x", TensorProto.FLOAT, x_shape)],
        outputs=[helper.make_tensor_value_info("output", TensorProto.FLOAT, x_shape)],
        value_info=[make_optional_tensor_value_info("optional", TensorProto.FLOAT, x_shape)],
    )
    model = helper.make_model(graph, producer_name="test_optional_get_element_tensor_ir")
    model.ir_version = 11
    model.opset_import[0].version = 18
    tvm_model = from_onnx(model, opset=18, keep_params_in_input=True)

    @I.ir_module
    class Expected:
        @R.function
        def main(x: R.Tensor((2, 3), dtype="float32")) -> R.Tensor((2, 3), dtype="float32"):
            R.func_attr({"num_input": 1})
            with R.dataflow():
                gv: R.Tensor((2, 3), dtype="float32") = x
                R.output(gv)
            return gv

    tvm.ir.assert_structural_equal(tvm_model, Expected)


def test_optional_get_element_sequence():
    seq_node, graph_inputs = construct_sequence(input_shape=[32, 32], num_tensors=4)
    index = make_constant_node("index", TensorProto.INT64, (), [1])
    optional_node = helper.make_node("Optional", ["sequence"], ["optional"])
    get_element_node = helper.make_node("OptionalGetElement", ["optional"], ["unwrapped"])
    sequence_at_node = helper.make_node("SequenceAt", ["unwrapped", "index"], ["output"])
    graph = helper.make_graph(
        [index, seq_node, optional_node, get_element_node, sequence_at_node],
        "test_optional_get_element_sequence",
        inputs=graph_inputs,
        outputs=[helper.make_tensor_value_info("output", TensorProto.FLOAT, [32, 32])],
        value_info=[make_optional_sequence_value_info("optional", TensorProto.FLOAT, [32, 32])],
    )
    model = helper.make_model(graph, producer_name="test_optional_get_element_sequence")
    model.ir_version = 11
    model.opset_import[0].version = 18
    tvm_model = from_onnx(model, opset=18, keep_params_in_input=True)

    @I.ir_module
    class Expected:
        @R.function
        def main(
            data0: R.Tensor((32, 32), dtype="float32"),
            data1: R.Tensor((32, 32), dtype="float32"),
            data2: R.Tensor((32, 32), dtype="float32"),
            data3: R.Tensor((32, 32), dtype="float32"),
        ) -> R.Tensor((32, 32), dtype="float32"):
            R.func_attr({"num_input": 4})
            with R.dataflow():
                gv: R.Tensor((32, 32), dtype="float32") = data1
                R.output(gv)
            return gv

    tvm.ir.assert_structural_equal(tvm_model, Expected)


def test_optional_without_input_requires_type_attr():
    tensor_type = helper.make_tensor_type_proto(TensorProto.FLOAT, [2, 3])
    optional_type = helper.make_optional_type_proto(tensor_type)
    optional_node = helper.make_node("Optional", [], ["optional"])
    graph = helper.make_graph(
        [optional_node],
        "test_optional_without_input_requires_type_attr",
        inputs=[],
        outputs=[helper.make_value_info("optional", optional_type)],
    )
    model = helper.make_model(graph, producer_name="test_optional_without_input_requires_type_attr")
    model.opset_import[0].version = 18

    with pytest.raises(ValueError, match="type attribute"):
        from_onnx(model, opset=18, keep_params_in_input=True)


def test_empty_optional_graph_output_raises():
    tensor_type = helper.make_tensor_type_proto(TensorProto.FLOAT, [2, 3])
    optional_type = helper.make_optional_type_proto(tensor_type)
    optional_node = helper.make_node("Optional", [], ["optional"], type=tensor_type)
    graph = helper.make_graph(
        [optional_node],
        "test_empty_optional_graph_output_raises",
        inputs=[],
        outputs=[helper.make_value_info("optional", optional_type)],
    )
    model = helper.make_model(graph, producer_name="test_empty_optional_graph_output_raises")
    model.opset_import[0].version = 18

    with pytest.raises(ValueError, match="Empty optional graph outputs are not supported"):
        from_onnx(model, opset=18, keep_params_in_input=True)


def test_optional_has_element_requires_one_input():
    has_element_node = helper.make_node("OptionalHasElement", [], ["output"])
    graph = helper.make_graph(
        [has_element_node],
        "test_optional_has_element_requires_one_input",
        inputs=[],
        outputs=[helper.make_tensor_value_info("output", TensorProto.BOOL, [])],
    )
    model = helper.make_model(graph, producer_name="test_optional_has_element_requires_one_input")
    model.opset_import[0].version = 18

    with pytest.raises(ValueError, match="expects one input"):
        from_onnx(model, opset=18, keep_params_in_input=True)


def test_optional_get_element_empty_raises():
    x_shape = [2, 3]
    tensor_type = helper.make_tensor_type_proto(TensorProto.FLOAT, x_shape)
    optional_type = helper.make_optional_type_proto(tensor_type)
    optional_node = helper.make_node("Optional", [], ["optional"], type=tensor_type)
    get_element_node = helper.make_node("OptionalGetElement", ["optional"], ["output"])
    graph = helper.make_graph(
        [optional_node, get_element_node],
        "test_optional_get_element_empty_raises",
        inputs=[],
        outputs=[helper.make_tensor_value_info("output", TensorProto.FLOAT, x_shape)],
        value_info=[helper.make_value_info("optional", optional_type)],
    )
    model = helper.make_model(graph, producer_name="test_optional_get_element_empty_raises")
    model.opset_import[0].version = 18
    with pytest.raises(ValueError, match="empty optional"):
        from_onnx(model, opset=18, keep_params_in_input=True)


def test_symbolic_shape_deduction():
    def verify_symbolic_shape_deduction(with_reshape_flatten, expected):
        index_node = helper.make_node(
            "Constant",
            inputs=[],
            outputs=["indices"],
            value=helper.make_tensor("indices", TensorProto.INT64, [], [0]),
        )
        shape_node = helper.make_node("Shape", ["data"], ["shape_output"])
        nodes = [index_node, shape_node]
        gather_input = "shape_output"

        if with_reshape_flatten:
            reshape_node = helper.make_node(
                "Reshape", ["shape_output", "target_shape"], ["reshaped_shape"]
            )
            nodes.append(reshape_node)
            gather_input = "reshaped_shape"

        gather_node = helper.make_node("Gather", [gather_input, "indices"], ["gather_output"])
        unsqueeze_node = helper.make_node(
            "Unsqueeze", ["gather_output", "axes"], ["unsqueeze_output"]
        )
        constant_of_shape_node = helper.make_node(
            "ConstantOfShape",
            ["unsqueeze_output"],
            ["output"],
            value=helper.make_tensor("value", TensorProto.FLOAT, [], [1]),
        )
        nodes.extend([gather_node, unsqueeze_node, constant_of_shape_node])

        initializers = [helper.make_tensor("axes", TensorProto.INT64, [1], vals=[0])]
        if with_reshape_flatten:
            initializers.append(
                helper.make_tensor("target_shape", TensorProto.INT64, [1], vals=[-1])
            )

        graph = helper.make_graph(
            nodes,
            "test_shape_deduction",
            inputs=[
                helper.make_tensor_value_info("data", TensorProto.FLOAT, ["batch", "seq"]),
            ],
            initializer=initializers,
            outputs=[helper.make_tensor_value_info("output", TensorProto.INT64, [1])],
        )
        model = helper.make_model(graph, producer_name="test_shape_deduction")
        tvm_model = from_onnx(model, keep_params_in_input=True)
        tvm.ir.assert_structural_equal(tvm_model["main"].without_attr("params"), expected["main"])

    @I.ir_module
    class ExpectedWithReshapeFlatten:
        @R.function
        def main(
            data: R.Tensor(("batch", "seq"), dtype="float32"),
            axes: R.Tensor((1,), dtype="int64"),
            target_shape: R.Tensor((1,), dtype="int64"),
        ) -> R.Tensor(("batch",), dtype="float32"):
            batch = T.int64()
            seq = T.int64()
            R.func_attr({"num_input": 1})
            with R.dataflow():
                gv: R.Tensor((batch,), dtype="float32") = R.broadcast_to(
                    R.const(1, "float32"), R.shape([batch])
                )
                R.output(gv)
            return gv

    @I.ir_module
    class ExpectedWithoutReshapeFlatten:
        @R.function
        def main(
            data: R.Tensor(("batch", "seq"), dtype="float32"),
            axes: R.Tensor((1,), dtype="int64"),
        ) -> R.Tensor(("batch",), dtype="float32"):
            batch = T.int64()
            seq = T.int64()
            R.func_attr({"num_input": 1})
            with R.dataflow():
                gv: R.Tensor((batch,), dtype="float32") = R.broadcast_to(
                    R.const(1, "float32"), R.shape([batch])
                )
                R.output(gv)
            return gv

    verify_symbolic_shape_deduction(False, ExpectedWithoutReshapeFlatten)
    verify_symbolic_shape_deduction(True, ExpectedWithReshapeFlatten)


def test_multi_inputs_with_same_symbolic_shape():
    concat_node = helper.make_node("Concat", ["data1", "data2"], ["output"], axis=1)

    graph = helper.make_graph(
        [concat_node],
        "test_multi_symbolic_shape_input",
        inputs=[
            helper.make_tensor_value_info("data1", TensorProto.FLOAT, ["batch", 1]),
            helper.make_tensor_value_info("data2", TensorProto.FLOAT, ["batch", 1]),
        ],
        outputs=[helper.make_tensor_value_info("output", TensorProto.FLOAT, ["batch", 2])],
    )
    model = helper.make_model(graph, producer_name="test_multi_symbolic_shape_input")
    tvm_model = from_onnx(model, keep_params_in_input=True)

    @I.ir_module
    class Expected:
        @R.function
        def main(
            data1: R.Tensor(("batch", 1), dtype="float32"),
            data2: R.Tensor(("batch", 1), dtype="float32"),
        ) -> R.Tensor(("batch", 2), dtype="float32"):
            batch = T.int64()
            R.func_attr({"num_input": 2})
            with R.dataflow():
                gv: R.Tensor((batch, 2), dtype="float32") = R.concat((data1, data2), axis=1)
                R.output(gv)
            return gv

    tvm.ir.assert_structural_equal(tvm_model, Expected)


def test_multi_ops_with_same_params():
    reshape_node_1 = helper.make_node("Reshape", ["a", "x"], ["b"])
    reshape_node_2 = helper.make_node("Reshape", ["b", "x"], ["c"])

    a_shape = [16]
    output_shape = [1, 16]

    graph = helper.make_graph(
        [reshape_node_1, reshape_node_2],
        "test_multi_ops_with_same_params",
        inputs=[
            helper.make_tensor_value_info("a", TensorProto.FLOAT, a_shape),
        ],
        initializer=[
            helper.make_tensor("x", TensorProto.INT64, [2], output_shape),
        ],
        outputs=[helper.make_tensor_value_info("c", TensorProto.FLOAT, output_shape)],
    )
    model = helper.make_model(graph, producer_name="test_multi_ops_with_same_params")
    tvm_model = from_onnx(model, keep_params_in_input=True)
    assert len(tvm_model["main"].attrs["params"]) == 1
    tvm_model["main"] = tvm_model["main"].without_attr("params")

    @I.ir_module
    class Expected:
        @R.function
        def main(
            a: R.Tensor((16,), dtype="float32"),
            v: R.Tensor((2,), dtype="int64"),
        ) -> R.Tensor((1, 16), dtype="float32"):
            R.func_attr({"num_input": 1})
            with R.dataflow():
                lv: R.Tensor((1, 16), dtype="float32") = R.reshape(a, R.shape([1, 16]))
                gv: R.Tensor((1, 16), dtype="float32") = R.reshape(lv, R.shape([1, 16]))
                R.output(gv)
            return gv

    tvm.ir.assert_structural_equal(tvm_model, Expected)


def test_params_names_start_with_onnx():
    reshape_node = helper.make_node("Reshape", ["a", "onnx::x"], ["b"])

    a_shape = [16]
    output_shape = [1, 16]

    graph = helper.make_graph(
        [reshape_node],
        "test_params_names_start_with_onnx",
        inputs=[
            helper.make_tensor_value_info("a", TensorProto.FLOAT, a_shape),
        ],
        initializer=[
            helper.make_tensor("onnx::x", TensorProto.INT64, [2], output_shape),
        ],
        outputs=[helper.make_tensor_value_info("b", TensorProto.FLOAT, output_shape)],
    )
    model = helper.make_model(graph, producer_name="test_params_names_start_with_onnx")
    tvm_model = from_onnx(model, keep_params_in_input=True)
    assert len(tvm_model["main"].attrs["params"]) == 1
    tvm_model["main"] = tvm_model["main"].without_attr("params")

    @I.ir_module
    class Expected:
        @R.function
        def main(
            a: R.Tensor((16,), dtype="float32"),
            v: R.Tensor((2,), dtype="int64"),
        ) -> R.Tensor((1, 16), dtype="float32"):
            R.func_attr({"num_input": 1})
            with R.dataflow():
                gv: R.Tensor((1, 16), dtype="float32") = R.reshape(a, R.shape([1, 16]))
                R.output(gv)
            return gv

    tvm.ir.assert_structural_equal(tvm_model, Expected)


def test_shape_dim_string_expression_graph_add():
    identity_node = helper.make_node("Identity", ["x"], ["y"])

    x_shape = ["A", "B", "A + B"]

    graph = helper.make_graph(
        [identity_node],
        "test_var_shape_dim_containing_expressions_onnx",
        inputs=[
            helper.make_tensor_value_info("x", TensorProto.FLOAT, x_shape),
        ],
        outputs=[helper.make_tensor_value_info("y", TensorProto.FLOAT, x_shape)],
    )
    model = helper.make_model(graph, producer_name="test_var_shape_dim_containing_expressions_onnx")

    tvm_model = from_onnx(model, opset=14, keep_params_in_input=True)

    # fmt: off
    @I.ir_module
    class Expected:
        @R.function
        def main(x: R.Tensor(("A", "B", "A + B"), dtype="float32")) -> R.Tensor(("A", "B", "A + B"), dtype="float32"):
            A = T.int64()
            B = T.int64()
            R.func_attr({"num_input": 1})
            with R.dataflow():
                gv: R.Tensor((A, B, A + B), dtype="float32") = x
                R.output(gv)
            return gv
    # fmt: on

    tvm.ir.assert_structural_equal(tvm_model, Expected)


def test_shape_dim_string_expression_graph_subtract():
    identity_node = helper.make_node("Identity", ["x"], ["y"])

    x_shape = ["A", "B", "A - B"]

    graph = helper.make_graph(
        [identity_node],
        "test_var_shape_dim_containing_expressions_onnx",
        inputs=[
            helper.make_tensor_value_info("x", TensorProto.FLOAT, x_shape),
        ],
        outputs=[helper.make_tensor_value_info("y", TensorProto.FLOAT, x_shape)],
    )
    model = helper.make_model(graph, producer_name="test_var_shape_dim_containing_expressions_onnx")

    tvm_model = from_onnx(model, opset=14, keep_params_in_input=True)

    # fmt: off
    @I.ir_module
    class Expected:
        @R.function
        def main(x: R.Tensor(("A", "B", "A - B"), dtype="float32")) -> R.Tensor(("A", "B", "A - B"), dtype="float32"):
            A = T.int64()
            B = T.int64()
            R.func_attr({"num_input": 1})
            with R.dataflow():
                gv: R.Tensor((A, B, A - B), dtype="float32") = x
                R.output(gv)
            return gv
    # fmt: on

    tvm.ir.assert_structural_equal(tvm_model, Expected)


def test_shape_dim_string_expression_graph_mul():
    identity_node = helper.make_node("Identity", ["x"], ["y"])

    x_shape = ["A", "B", "A * B"]

    graph = helper.make_graph(
        [identity_node],
        "test_var_shape_dim_containing_expressions_onnx",
        inputs=[
            helper.make_tensor_value_info("x", TensorProto.FLOAT, x_shape),
        ],
        outputs=[helper.make_tensor_value_info("y", TensorProto.FLOAT, x_shape)],
    )
    model = helper.make_model(graph, producer_name="test_var_shape_dim_containing_expressions_onnx")

    tvm_model = from_onnx(model, opset=14, keep_params_in_input=True)

    # fmt: off
    @I.ir_module
    class Expected:
        @R.function
        def main(x: R.Tensor(("A", "B", "A * B"), dtype="float32")) -> R.Tensor(("A", "B", "A * B"), dtype="float32"):
            A = T.int64()
            B = T.int64()
            R.func_attr({"num_input": 1})
            with R.dataflow():
                gv: R.Tensor((A, B, A * B), dtype="float32") = x
                R.output(gv)
            return gv
    # fmt: on

    tvm.ir.assert_structural_equal(tvm_model, Expected)


def test_shape_dim_string_expression_graph_div_1():
    identity_node = helper.make_node("Identity", ["x"], ["y"])

    # this will result in a floordiv despite not using // since the operands are always int
    x_shape = ["A", "B", "A / B"]

    graph = helper.make_graph(
        [identity_node],
        "test_var_shape_dim_containing_expressions_onnx",
        inputs=[
            helper.make_tensor_value_info("x", TensorProto.FLOAT, x_shape),
        ],
        outputs=[helper.make_tensor_value_info("y", TensorProto.FLOAT, x_shape)],
    )
    model = helper.make_model(graph, producer_name="test_var_shape_dim_containing_expressions_onnx")

    tvm_model = from_onnx(model, opset=14, keep_params_in_input=True)

    # fmt: off
    @I.ir_module
    class Expected:
        @R.function
        def main(x: R.Tensor(("A", "B", "A // B"), dtype="float32")) -> R.Tensor(("A", "B", "A // B"), dtype="float32"):
            A = T.int64()
            B = T.int64()
            R.func_attr({"num_input": 1})
            with R.dataflow():
                gv: R.Tensor((A, B, A // B), dtype="float32") = x
                R.output(gv)
            return gv

    # fmt: on

    tvm.ir.assert_structural_equal(tvm_model, Expected)


def test_shape_dim_string_expression_graph_div_2():
    identity_node = helper.make_node("Identity", ["x"], ["y"])

    x_shape = ["A", "B", "A // B"]

    graph = helper.make_graph(
        [identity_node],
        "test_var_shape_dim_containing_expressions_onnx",
        inputs=[
            helper.make_tensor_value_info("x", TensorProto.FLOAT, x_shape),
        ],
        outputs=[helper.make_tensor_value_info("y", TensorProto.FLOAT, x_shape)],
    )
    model = helper.make_model(graph, producer_name="test_var_shape_dim_containing_expressions_onnx")

    tvm_model = from_onnx(model, opset=14, keep_params_in_input=True)

    # fmt: off
    @I.ir_module
    class Expected:
        @R.function
        def main(x: R.Tensor(("A", "B", "A // B"), dtype="float32")) -> R.Tensor(("A", "B", "A // B"), dtype="float32"):
            A = T.int64()
            B = T.int64()
            R.func_attr({"num_input": 1})
            with R.dataflow():
                gv: R.Tensor((A, B, A // B), dtype="float32") = x
                R.output(gv)
            return gv
    # fmt: on

    tvm.ir.assert_structural_equal(tvm_model, Expected)


@I.ir_module
class ExpectedNMSFiveBoxes:
    @R.function
    def main(
        boxes: R.Tensor((1, 5, 4), dtype="float32"),
        scores: R.Tensor((1, 2, 5), dtype="float32"),
        max_output_boxes_per_class: R.Tensor((1,), dtype="int64"),
        iou_threshold: R.Tensor((1,), dtype="float32"),
        score_threshold: R.Tensor((1,), dtype="float32"),
    ):
        R.func_attr({"num_input": 2})
        with R.dataflow():
            lv = R.vision.all_class_non_max_suppression(
                boxes,
                scores,
                R.const(3, "int64"),
                R.const(0.5, "float32"),
                R.const(0.10000000149011612, "float32"),
                "onnx",
            )
            lv1 = lv[0]
            gv = lv1
            R.output(gv)
        return gv


@I.ir_module
class ExpectedNMSFourBoxesDefaultParams:
    @R.function
    def main(
        boxes: R.Tensor((1, 4, 4), dtype="float32"),
        scores: R.Tensor((1, 1, 4), dtype="float32"),
    ):
        R.func_attr({"num_input": 2})
        with R.dataflow():
            lv = R.vision.all_class_non_max_suppression(
                boxes,
                scores,
                R.const(0, "int64"),
                R.const(0.5, "float32"),
                R.const(0.0, "float32"),
                "onnx",
            )
            lv1 = lv[0]
            gv = lv1
            R.output(gv)
        return gv


@I.ir_module
class ExpectedNMSFourBoxesWithMaxParam:
    @R.function
    def main(
        boxes: R.Tensor((1, 4, 4), dtype="float32"),
        scores: R.Tensor((1, 1, 4), dtype="float32"),
        max_output_boxes_per_class: R.Tensor((1,), dtype="int64"),
    ):
        R.func_attr({"num_input": 2})
        with R.dataflow():
            lv = R.vision.all_class_non_max_suppression(
                boxes,
                scores,
                R.const(0, "int64"),
                R.const(0.5, "float32"),
                R.const(0.0, "float32"),
                "onnx",
            )
            lv1 = lv[0]
            gv = lv1
            R.output(gv)
        return gv


@I.ir_module
class ExpectedNMSFourBoxes:
    @R.function
    def main(
        boxes: R.Tensor((1, 4, 4), dtype="float32"),
        scores: R.Tensor((1, 1, 4), dtype="float32"),
        max_output_boxes_per_class: R.Tensor((1,), dtype="int64"),
        iou_threshold: R.Tensor((1,), dtype="float32"),
        score_threshold: R.Tensor((1,), dtype="float32"),
    ):
        R.func_attr({"num_input": 2})
        with R.dataflow():
            lv = R.vision.all_class_non_max_suppression(
                boxes,
                scores,
                R.const(2, "int64"),
                R.const(0.10000000149011612, "float32"),
                R.const(0.10000000149011612, "float32"),
                "onnx",
            )
            lv1 = lv[0]
            gv = lv1
            R.output(gv)
        return gv


@I.ir_module
class ExpectedNMSThreeBoxesTwoClasses:
    @R.function
    def main(
        boxes: R.Tensor((1, 3, 4), dtype="float32"),
        scores: R.Tensor((1, 2, 3), dtype="float32"),
        max_output_boxes_per_class: R.Tensor((1,), dtype="int64"),
        iou_threshold: R.Tensor((1,), dtype="float32"),
        score_threshold: R.Tensor((1,), dtype="float32"),
    ):
        R.func_attr({"num_input": 2})
        with R.dataflow():
            lv = R.vision.all_class_non_max_suppression(
                boxes,
                scores,
                R.const(2, "int64"),
                R.const(0.5, "float32"),
                R.const(0.10000000149011612, "float32"),
                "onnx",
            )
            lv1 = lv[0]
            gv = lv1
            R.output(gv)
        return gv


@I.ir_module
class ExpectedNMSThreeBoxesOneClass:
    @R.function
    def main(
        boxes: R.Tensor((1, 3, 4), dtype="float32"),
        scores: R.Tensor((1, 1, 3), dtype="float32"),
        max_output_boxes_per_class: R.Tensor((1,), dtype="int64"),
        iou_threshold: R.Tensor((1,), dtype="float32"),
        score_threshold: R.Tensor((1,), dtype="float32"),
    ):
        R.func_attr({"num_input": 2})
        with R.dataflow():
            lv = R.vision.all_class_non_max_suppression(
                boxes,
                scores,
                R.const(2, "int64"),
                R.const(0.5, "float32"),
                R.const(0.10000000149011612, "float32"),
                "onnx",
            )
            lv1 = lv[0]
            gv = lv1
            R.output(gv)
        return gv


@I.ir_module
class ExpectedNMSThreeBoxesOneClassScoreThreshold:
    @R.function
    def main(
        boxes: R.Tensor((1, 3, 4), dtype="float32"),
        scores: R.Tensor((1, 1, 3), dtype="float32"),
        max_output_boxes_per_class: R.Tensor((1,), dtype="int64"),
        iou_threshold: R.Tensor((1,), dtype="float32"),
        score_threshold: R.Tensor((1,), dtype="float32"),
    ):
        R.func_attr({"num_input": 2})
        with R.dataflow():
            lv = R.vision.all_class_non_max_suppression(
                boxes,
                scores,
                R.const(3, "int64"),
                R.const(0.10000000149011612, "float32"),
                R.const(0.05000000074505806, "float32"),
                "onnx",
            )
            lv1 = lv[0]
            gv = lv1
            R.output(gv)
        return gv


def _assert_nms_import(
    model,
    boxes_shape,
    scores_shape,
    expected,
    center_point_box=0,
    nms_params=None,
):
    assert center_point_box == 0
    nms_params = nms_params or []

    tvm_model = from_onnx(model, opset=11, keep_params_in_input=True)
    if nms_params:
        assert len(tvm_model["main"].attrs["params"]) == len(nms_params)
        tvm_model["main"] = tvm_model["main"].without_attr("params")

    tvm.ir.assert_structural_equal(tvm_model, expected)


def test_nms():
    """NonMaxSuppression should import as all_class_non_max_suppression."""
    nms_node = helper.make_node(
        "NonMaxSuppression",
        ["boxes", "scores", "max_output_boxes_per_class", "iou_threshold", "score_threshold"],
        ["selected_indices"],
        center_point_box=0,
    )

    boxes_shape = [1, 5, 4]  # batch_size, num_boxes, 4
    scores_shape = [1, 2, 5]  # batch_size, num_classes, num_boxes

    graph = helper.make_graph(
        [nms_node],
        "nms_test",
        inputs=[
            helper.make_tensor_value_info("boxes", TensorProto.FLOAT, boxes_shape),
            helper.make_tensor_value_info("scores", TensorProto.FLOAT, scores_shape),
        ],
        initializer=[
            helper.make_tensor("max_output_boxes_per_class", TensorProto.INT64, [1], [3]),
            helper.make_tensor("iou_threshold", TensorProto.FLOAT, [1], [0.5]),
            helper.make_tensor("score_threshold", TensorProto.FLOAT, [1], [0.1]),
        ],
        outputs=[helper.make_tensor_value_info("selected_indices", TensorProto.INT64, [0, 3])],
    )

    model = helper.make_model(graph, producer_name="nms_test")
    model.ir_version = 8
    model.opset_import[0].version = 11

    _assert_nms_import(
        model,
        boxes_shape,
        scores_shape,
        ExpectedNMSFiveBoxes,
        nms_params=[
            ("max_output_boxes_per_class", [1], "int64", 3),
            ("iou_threshold", [1], "float32", 0.5),
            ("score_threshold", [1], "float32", 0.1),
        ],
    )


def test_nms_scalar_shape1_constants():
    """Scalar params given as 1-D single-element constants must import (NumPy 2.x cast)."""
    nms_node = helper.make_node(
        "NonMaxSuppression",
        ["boxes", "scores", "max_output_boxes_per_class", "iou_threshold", "score_threshold"],
        ["selected_indices"],
    )
    graph = helper.make_graph(
        [nms_node],
        "nms_scalar_shape1",
        inputs=[
            helper.make_tensor_value_info("boxes", TensorProto.FLOAT, [1, 5, 4]),
            helper.make_tensor_value_info("scores", TensorProto.FLOAT, [1, 1, 5]),
        ],
        initializer=[
            helper.make_tensor("max_output_boxes_per_class", TensorProto.INT64, [1], [3]),
            helper.make_tensor("iou_threshold", TensorProto.FLOAT, [1], [0.5]),
            helper.make_tensor("score_threshold", TensorProto.FLOAT, [1], [0.0]),
        ],
        outputs=[helper.make_tensor_value_info("selected_indices", TensorProto.INT64, [0, 3])],
    )
    model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 18)])
    # Default import folds initializers to relax.Constant, exercising the scalar-cast path.
    from_onnx(model)


def test_nms_max_output_boxes_per_class_zero():
    """ONNX default for max_output_boxes_per_class should import as 0."""

    def verify(with_explicit_max, expected):
        node_inputs = ["boxes", "scores"]
        initializer = []
        nms_params = None
        if with_explicit_max:
            node_inputs.append("max_output_boxes_per_class")
            initializer.append(
                helper.make_tensor("max_output_boxes_per_class", TensorProto.INT64, [1], [0])
            )
            nms_params = [("max_output_boxes_per_class", [1], "int64", 0)]

        nms_node = helper.make_node(
            "NonMaxSuppression",
            node_inputs,
            ["selected_indices"],
            center_point_box=0,
        )

        boxes_shape = [1, 4, 4]
        scores_shape = [1, 1, 4]
        graph = helper.make_graph(
            [nms_node],
            "nms_max_output_boxes_per_class_zero",
            inputs=[
                helper.make_tensor_value_info("boxes", TensorProto.FLOAT, boxes_shape),
                helper.make_tensor_value_info("scores", TensorProto.FLOAT, scores_shape),
            ],
            initializer=initializer,
            outputs=[helper.make_tensor_value_info("selected_indices", TensorProto.INT64, [0, 3])],
        )

        model = helper.make_model(graph, producer_name="nms_max_output_boxes_per_class_zero")
        model.ir_version = 8
        model.opset_import[0].version = 11

        _assert_nms_import(
            model,
            boxes_shape,
            scores_shape,
            expected,
            nms_params=nms_params,
        )

    verify(False, ExpectedNMSFourBoxesDefaultParams)
    verify(True, ExpectedNMSFourBoxesWithMaxParam)


def test_nms_algorithm_correctness():
    """NMS import should pass max boxes, IoU, and score threshold constants."""
    nms_node = helper.make_node(
        "NonMaxSuppression",
        ["boxes", "scores", "max_output_boxes_per_class", "iou_threshold", "score_threshold"],
        ["selected_indices"],
        center_point_box=0,
    )

    boxes_shape = [1, 3, 4]  # batch_size, num_boxes, 4
    scores_shape = [1, 2, 3]  # batch_size, num_classes, num_boxes

    graph = helper.make_graph(
        [nms_node],
        "nms_test_correctness",
        inputs=[
            helper.make_tensor_value_info("boxes", TensorProto.FLOAT, boxes_shape),
            helper.make_tensor_value_info("scores", TensorProto.FLOAT, scores_shape),
        ],
        initializer=[
            helper.make_tensor(
                "max_output_boxes_per_class", TensorProto.INT64, [1], [2]
            ),  # Only 2 boxes per class
            helper.make_tensor("iou_threshold", TensorProto.FLOAT, [1], [0.5]),  # IoU threshold 0.5
            helper.make_tensor(
                "score_threshold", TensorProto.FLOAT, [1], [0.1]
            ),  # Score threshold 0.1
        ],
        outputs=[helper.make_tensor_value_info("selected_indices", TensorProto.INT64, [4, 3])],
    )

    model = helper.make_model(graph, producer_name="nms_test_correctness")

    _assert_nms_import(
        model,
        boxes_shape,
        scores_shape,
        ExpectedNMSThreeBoxesTwoClasses,
        nms_params=[
            ("max_output_boxes_per_class", [1], "int64", 2),
            ("iou_threshold", [1], "float32", 0.5),
            ("score_threshold", [1], "float32", 0.1),
        ],
    )


def test_nms_iou_suppression():
    """NMS import should pass the IoU threshold constant."""
    nms_node = helper.make_node(
        "NonMaxSuppression",
        ["boxes", "scores", "max_output_boxes_per_class", "iou_threshold", "score_threshold"],
        ["selected_indices"],
        center_point_box=0,
    )

    boxes_shape = [1, 3, 4]
    scores_shape = [1, 1, 3]

    graph = helper.make_graph(
        [nms_node],
        "nms_test_iou_suppression",
        inputs=[
            helper.make_tensor_value_info("boxes", TensorProto.FLOAT, boxes_shape),
            helper.make_tensor_value_info("scores", TensorProto.FLOAT, scores_shape),
        ],
        initializer=[
            helper.make_tensor("max_output_boxes_per_class", TensorProto.INT64, [1], [2]),
            helper.make_tensor("iou_threshold", TensorProto.FLOAT, [1], [0.5]),  # IoU threshold 0.5
            helper.make_tensor("score_threshold", TensorProto.FLOAT, [1], [0.1]),
        ],
        outputs=[helper.make_tensor_value_info("selected_indices", TensorProto.INT64, [2, 3])],
    )

    model = helper.make_model(graph, producer_name="nms_test_iou_suppression")
    model.ir_version = 8
    model.opset_import[0].version = 11

    _assert_nms_import(
        model,
        boxes_shape,
        scores_shape,
        ExpectedNMSThreeBoxesOneClass,
        nms_params=[
            ("max_output_boxes_per_class", [1], "int64", 2),
            ("iou_threshold", [1], "float32", 0.5),
            ("score_threshold", [1], "float32", 0.1),
        ],
    )


def test_nms_max_boxes_limit():
    """NMS import should pass max_output_boxes_per_class."""
    nms_node = helper.make_node(
        "NonMaxSuppression",
        ["boxes", "scores", "max_output_boxes_per_class", "iou_threshold", "score_threshold"],
        ["selected_indices"],
        center_point_box=0,
    )

    boxes_shape = [1, 4, 4]
    scores_shape = [1, 1, 4]

    graph = helper.make_graph(
        [nms_node],
        "nms_test_max_boxes_limit",
        inputs=[
            helper.make_tensor_value_info("boxes", TensorProto.FLOAT, boxes_shape),
            helper.make_tensor_value_info("scores", TensorProto.FLOAT, scores_shape),
        ],
        initializer=[
            helper.make_tensor(
                "max_output_boxes_per_class", TensorProto.INT64, [1], [2]
            ),  # Limit to 2 boxes
            helper.make_tensor("iou_threshold", TensorProto.FLOAT, [1], [0.1]),  # Low IoU threshold
            helper.make_tensor("score_threshold", TensorProto.FLOAT, [1], [0.1]),
        ],
        outputs=[helper.make_tensor_value_info("selected_indices", TensorProto.INT64, [2, 3])],
    )

    model = helper.make_model(graph, producer_name="nms_test_max_boxes_limit")
    model.ir_version = 8
    model.opset_import[0].version = 11

    _assert_nms_import(
        model,
        boxes_shape,
        scores_shape,
        ExpectedNMSFourBoxes,
        nms_params=[
            ("max_output_boxes_per_class", [1], "int64", 2),
            ("iou_threshold", [1], "float32", 0.1),
            ("score_threshold", [1], "float32", 0.1),
        ],
    )


def test_nms_score_threshold():
    """NMS import should pass the score threshold constant."""
    nms_node = helper.make_node(
        "NonMaxSuppression",
        ["boxes", "scores", "max_output_boxes_per_class", "iou_threshold", "score_threshold"],
        ["selected_indices"],
        center_point_box=0,
    )

    boxes_shape = [1, 3, 4]
    scores_shape = [1, 1, 3]

    graph = helper.make_graph(
        [nms_node],
        "nms_test_score_threshold",
        inputs=[
            helper.make_tensor_value_info("boxes", TensorProto.FLOAT, boxes_shape),
            helper.make_tensor_value_info("scores", TensorProto.FLOAT, scores_shape),
        ],
        initializer=[
            helper.make_tensor("max_output_boxes_per_class", TensorProto.INT64, [1], [3]),
            helper.make_tensor("iou_threshold", TensorProto.FLOAT, [1], [0.1]),
            helper.make_tensor("score_threshold", TensorProto.FLOAT, [1], [0.05]),
        ],
        outputs=[helper.make_tensor_value_info("selected_indices", TensorProto.INT64, [3, 3])],
    )

    model = helper.make_model(graph, producer_name="nms_test_score_threshold")
    model.ir_version = 8
    model.opset_import[0].version = 11

    _assert_nms_import(
        model,
        boxes_shape,
        scores_shape,
        ExpectedNMSThreeBoxesOneClassScoreThreshold,
        nms_params=[
            ("max_output_boxes_per_class", [1], "int64", 3),
            ("iou_threshold", [1], "float32", 0.1),
            ("score_threshold", [1], "float32", 0.05),
        ],
    )


# align_corners=None omits the attribute, exercising the ONNX default of 0.
def test_affine_grid():
    def verify_affine_grid(align_corners, expected):
        attrs = {} if align_corners is None else {"align_corners": align_corners}
        affine_grid_node = helper.make_node("AffineGrid", ["theta", "size"], ["grid"], **attrs)

        graph = helper.make_graph(
            [affine_grid_node],
            "affine_grid_test",
            inputs=[
                helper.make_tensor_value_info("theta", TensorProto.FLOAT, [2, 2, 3]),
            ],
            initializer=[
                helper.make_tensor("size", TensorProto.INT64, [4], [2, 3, 16, 16]),
            ],
            outputs=[
                helper.make_tensor_value_info("grid", TensorProto.FLOAT, [2, 16, 16, 2]),
            ],
        )

        model = helper.make_model(
            graph, producer_name="affine_grid_test", opset_imports=[helper.make_opsetid("", 20)]
        )
        tvm_model = from_onnx(model, opset=20, keep_params_in_input=True)
        assert len(tvm_model["main"].attrs["params"]) == 1
        tvm_model["main"] = tvm_model["main"].without_attr("params")
        tvm.ir.assert_structural_equal(tvm_model, expected)

    @I.ir_module
    class ExpectedAlignCorners:
        @R.function
        def main(
            theta: R.Tensor((2, 2, 3), dtype="float32"),
            size: R.Tensor((4,), dtype="int64"),
        ) -> R.Tensor((2, 16, 16, 2), dtype="float32"):
            R.func_attr({"num_input": 1})
            with R.dataflow():
                lv: R.Tensor((2, 2, 16, 16), dtype="float32") = R.image.affine_grid(
                    theta, size=(16, 16), align_corners=True
                )
                lv1: R.Tensor((2, 16, 16, 2), dtype="float32") = R.permute_dims(
                    lv, axes=[0, 2, 3, 1]
                )
                gv: R.Tensor((2, 16, 16, 2), dtype="float32") = lv1
                R.output(gv)
            return gv

    @I.ir_module
    class ExpectedDefaultAlignCorners:
        @R.function
        def main(
            theta: R.Tensor((2, 2, 3), dtype="float32"),
            size: R.Tensor((4,), dtype="int64"),
        ) -> R.Tensor((2, 16, 16, 2), dtype="float32"):
            R.func_attr({"num_input": 1})
            with R.dataflow():
                lv: R.Tensor((2, 2, 16, 16), dtype="float32") = R.image.affine_grid(
                    theta, size=(16, 16), align_corners=False
                )
                lv1: R.Tensor((2, 16, 16, 2), dtype="float32") = R.permute_dims(
                    lv, axes=[0, 2, 3, 1]
                )
                gv: R.Tensor((2, 16, 16, 2), dtype="float32") = lv1
                R.output(gv)
            return gv

    verify_affine_grid(None, ExpectedDefaultAlignCorners)
    verify_affine_grid(0, ExpectedDefaultAlignCorners)
    verify_affine_grid(1, ExpectedAlignCorners)


def test_affine_grid_3d():
    affine_grid_node = helper.make_node(
        "AffineGrid",
        ["theta", "size"],
        ["grid"],
        align_corners=1,
    )

    graph = helper.make_graph(
        [affine_grid_node],
        "affine_grid_3d_test",
        inputs=[
            helper.make_tensor_value_info("theta", TensorProto.FLOAT, [2, 3, 4]),
        ],
        initializer=[
            helper.make_tensor("size", TensorProto.INT64, [5], [2, 3, 8, 16, 16]),
        ],
        outputs=[
            helper.make_tensor_value_info("grid", TensorProto.FLOAT, [2, 8, 16, 16, 3]),
        ],
    )

    model = helper.make_model(graph, producer_name="affine_grid_3d_test")

    tvm_model = from_onnx(model, opset=20, keep_params_in_input=True)
    assert len(tvm_model["main"].attrs["params"]) == 1
    tvm_model["main"] = tvm_model["main"].without_attr("params")

    @I.ir_module
    class Expected:
        @R.function
        def main(
            theta: R.Tensor((2, 3, 4), dtype="float32"),
            size: R.Tensor((5,), dtype="int64"),
        ) -> R.Tensor((2, 8, 16, 16, 3), dtype="float32"):
            R.func_attr({"num_input": 1})
            with R.dataflow():
                lv: R.Tensor((2, 3, 8, 16, 16), dtype="float32") = R.image.affine_grid(
                    theta, size=(8, 16, 16), align_corners=True
                )
                lv1: R.Tensor((2, 8, 16, 16, 3), dtype="float32") = R.permute_dims(
                    lv, axes=[0, 2, 3, 4, 1]
                )
                gv: R.Tensor((2, 8, 16, 16, 3), dtype="float32") = lv1
                R.output(gv)
            return gv

    tvm.ir.assert_structural_equal(tvm_model, Expected)


@pytest.mark.parametrize("mode", ["bilinear", "nearest", "bicubic"])
@pytest.mark.parametrize("padding_mode", ["zeros", "border", "reflection"])
@pytest.mark.parametrize("align_corners", [0, 1])
def test_grid_sample(mode, padding_mode, align_corners):
    x_shape = [1, 3, 4, 4]
    grid_shape = [1, 2, 2, 2]
    out_shape = [x_shape[0], x_shape[1], grid_shape[1], grid_shape[2]]

    node = helper.make_node(
        "GridSample",
        inputs=["X", "grid"],
        outputs=["Y"],
        mode=mode,
        padding_mode=padding_mode,
        align_corners=align_corners,
    )

    graph = helper.make_graph(
        [node],
        "grid_sample_test",
        inputs=[
            helper.make_tensor_value_info("X", TensorProto.FLOAT, x_shape),
            helper.make_tensor_value_info("grid", TensorProto.FLOAT, grid_shape),
        ],
        outputs=[
            helper.make_tensor_value_info("Y", TensorProto.FLOAT, out_shape),
        ],
    )

    model = helper.make_model(
        graph, producer_name="grid_sample_test", opset_imports=[helper.make_opsetid("", 16)]
    )
    tvm_model = from_onnx(model, opset=16, keep_params_in_input=True)

    @I.ir_module
    class ExpectedGridSample4D:
        @R.function
        def main(
            X: R.Tensor((1, 3, 4, 4), dtype="float32"),
            grid: R.Tensor((1, 2, 2, 2), dtype="float32"),
        ) -> R.Tensor((1, 3, 2, 2), dtype="float32"):
            R.func_attr({"num_input": 2})
            with R.dataflow():
                lv: R.Tensor((1, 2, 2, 2), dtype="float32") = R.permute_dims(
                    grid, axes=[0, 3, 1, 2]
                )
                gv: R.Tensor((1, 3, 2, 2), dtype="float32") = R.image.grid_sample(
                    X,
                    lv,
                    method=mode,
                    layout="NCHW",
                    padding_mode=padding_mode,
                    align_corners=bool(align_corners),
                )
                R.output(gv)
            return gv

    tvm.ir.assert_structural_equal(tvm_model, ExpectedGridSample4D)


@pytest.mark.parametrize("mode", ["bilinear", "nearest"])
@pytest.mark.parametrize("padding_mode", ["zeros", "border", "reflection"])
@pytest.mark.parametrize("align_corners", [0, 1])
def test_grid_sample_5d(mode, padding_mode, align_corners):
    x_shape = [1, 1, 4, 4, 4]
    grid_shape = [1, 4, 4, 4, 3]
    out_shape = [x_shape[0], x_shape[1], grid_shape[1], grid_shape[2], grid_shape[3]]

    node = helper.make_node(
        "GridSample",
        inputs=["X", "grid"],
        outputs=["Y"],
        mode=mode,
        padding_mode=padding_mode,
        align_corners=align_corners,
    )

    graph = helper.make_graph(
        [node],
        "grid_sample_5d_test",
        inputs=[
            helper.make_tensor_value_info("X", TensorProto.FLOAT, x_shape),
            helper.make_tensor_value_info("grid", TensorProto.FLOAT, grid_shape),
        ],
        outputs=[
            helper.make_tensor_value_info("Y", TensorProto.FLOAT, out_shape),
        ],
    )

    model = helper.make_model(
        graph, producer_name="grid_sample_5d_test", opset_imports=[helper.make_opsetid("", 16)]
    )
    tvm_model = from_onnx(model, opset=16, keep_params_in_input=True)

    @I.ir_module
    class ExpectedGridSample5D:
        @R.function
        def main(
            X: R.Tensor((1, 1, 4, 4, 4), dtype="float32"),
            grid: R.Tensor((1, 4, 4, 4, 3), dtype="float32"),
        ) -> R.Tensor((1, 1, 4, 4, 4), dtype="float32"):
            R.func_attr({"num_input": 2})
            with R.dataflow():
                lv: R.Tensor((1, 3, 4, 4, 4), dtype="float32") = R.permute_dims(
                    grid, axes=[0, 4, 1, 2, 3]
                )
                gv: R.Tensor((1, 1, 4, 4, 4), dtype="float32") = R.image.grid_sample(
                    X,
                    lv,
                    method=mode,
                    layout="NCDHW",
                    padding_mode=padding_mode,
                    align_corners=bool(align_corners),
                )
                R.output(gv)
            return gv

    tvm.ir.assert_structural_equal(tvm_model, ExpectedGridSample5D)


def test_grid_sample_5d_cubic_unsupported():
    x_shape = [1, 1, 4, 4, 4]
    grid_shape = [1, 2, 3, 5, 3]
    out_shape = [x_shape[0], x_shape[1], grid_shape[1], grid_shape[2], grid_shape[3]]

    node = helper.make_node(
        "GridSample",
        inputs=["X", "grid"],
        outputs=["Y"],
        mode="cubic",
    )

    graph = helper.make_graph(
        [node],
        "grid_sample_5d_cubic_unsupported_test",
        inputs=[
            helper.make_tensor_value_info("X", TensorProto.FLOAT, x_shape),
            helper.make_tensor_value_info("grid", TensorProto.FLOAT, grid_shape),
        ],
        outputs=[
            helper.make_tensor_value_info("Y", TensorProto.FLOAT, out_shape),
        ],
    )

    model = helper.make_model(graph, producer_name="grid_sample_5d_cubic_unsupported_test")
    with pytest.raises(
        NotImplementedError,
        match="5D .*GridSample with mode='cubic' is not supported",
    ):
        from_onnx(model, opset=16, keep_params_in_input=True)


def test_grid_sample_4d_non_square_output_shape():
    x_shape = [1, 3, 4, 4]
    grid_shape = [1, 3, 5, 2]
    out_shape = [x_shape[0], x_shape[1], grid_shape[1], grid_shape[2]]

    node = helper.make_node(
        "GridSample",
        inputs=["X", "grid"],
        outputs=["Y"],
        mode="bilinear",
    )

    graph = helper.make_graph(
        [node],
        "grid_sample_4d_non_square_output_shape_test",
        inputs=[
            helper.make_tensor_value_info("X", TensorProto.FLOAT, x_shape),
            helper.make_tensor_value_info("grid", TensorProto.FLOAT, grid_shape),
        ],
        outputs=[
            helper.make_tensor_value_info("Y", TensorProto.FLOAT, out_shape),
        ],
    )

    model = helper.make_model(graph, producer_name="grid_sample_4d_non_square_output_shape_test")
    tvm_model = from_onnx(model, opset=16, keep_params_in_input=True)

    @I.ir_module
    class Expected:
        @R.function
        def main(
            X: R.Tensor((1, 3, 4, 4), dtype="float32"),
            grid: R.Tensor((1, 3, 5, 2), dtype="float32"),
        ) -> R.Tensor((1, 3, 3, 5), dtype="float32"):
            R.func_attr({"num_input": 2})
            with R.dataflow():
                lv: R.Tensor((1, 2, 3, 5), dtype="float32") = R.permute_dims(
                    grid, axes=[0, 3, 1, 2]
                )
                gv: R.Tensor((1, 3, 3, 5), dtype="float32") = R.image.grid_sample(
                    X,
                    lv,
                    method="bilinear",
                    layout="NCHW",
                    padding_mode="zeros",
                    align_corners=False,
                )
                R.output(gv)
            return gv

    tvm.ir.assert_structural_equal(tvm_model, Expected)


def test_grid_sample_unsupported_rank():
    x_shape = [1, 3, 4]
    grid_shape = [1, 4, 2]

    node = helper.make_node(
        "GridSample",
        inputs=["X", "grid"],
        outputs=["Y"],
        mode="bilinear",
    )

    graph = helper.make_graph(
        [node],
        "grid_sample_unsupported_rank_test",
        inputs=[
            helper.make_tensor_value_info("X", TensorProto.FLOAT, x_shape),
            helper.make_tensor_value_info("grid", TensorProto.FLOAT, grid_shape),
        ],
        outputs=[
            helper.make_tensor_value_info("Y", TensorProto.FLOAT, x_shape),
        ],
    )

    model = helper.make_model(graph, producer_name="grid_sample_unsupported_rank_test")
    with pytest.raises(NotImplementedError, match="GridSample only supports 4D or 5D input"):
        from_onnx(model, opset=16, keep_params_in_input=True)


def test_grid_sample_linear_mode_translation():
    """Test that ONNX mode='linear' is correctly translated to 'bilinear'.

    The ONNX spec defines 'linear' as a valid mode for GridSample, but
    onnxruntime rejects it in practice. Real ONNX models exported from
    frameworks like PyTorch may still use 'linear'. We verify the translation
    by inspecting the Relax IR directly rather than running check_correctness.
    """
    x_shape = [1, 3, 4, 4]
    grid_shape = [1, 2, 2, 2]

    node = helper.make_node(
        "GridSample",
        inputs=["X", "grid"],
        outputs=["Y"],
        mode="linear",
    )

    graph = helper.make_graph(
        [node],
        "grid_sample_linear_test",
        inputs=[
            helper.make_tensor_value_info("X", TensorProto.FLOAT, x_shape),
            helper.make_tensor_value_info("grid", TensorProto.FLOAT, grid_shape),
        ],
        outputs=[
            helper.make_tensor_value_info(
                "Y", TensorProto.FLOAT, [x_shape[0], x_shape[1], grid_shape[1], grid_shape[2]]
            ),
        ],
    )

    model = helper.make_model(graph, producer_name="grid_sample_linear_test")
    tvm_model = from_onnx(model, opset=16, keep_params_in_input=True)

    @I.ir_module
    class Expected:
        @R.function
        def main(
            X: R.Tensor((1, 3, 4, 4), dtype="float32"),
            grid: R.Tensor((1, 2, 2, 2), dtype="float32"),
        ) -> R.Tensor((1, 3, 2, 2), dtype="float32"):
            R.func_attr({"num_input": 2})
            with R.dataflow():
                lv: R.Tensor((1, 2, 2, 2), dtype="float32") = R.permute_dims(
                    grid, axes=[0, 3, 1, 2]
                )
                gv: R.Tensor((1, 3, 2, 2), dtype="float32") = R.image.grid_sample(
                    X,
                    lv,
                    method="bilinear",
                    layout="NCHW",
                    padding_mode="zeros",
                    align_corners=False,
                )
                R.output(gv)
            return gv

    tvm.ir.assert_structural_equal(tvm_model, Expected)


def test_grid_sample_cubic_mode_translation():
    """Test that ONNX mode='cubic' is correctly translated to 'bicubic'.

    The ONNX spec defines 'cubic' as a valid mode for GridSample, but
    TVM uses 'bicubic'. We verify the translation by inspecting the
    Relax IR directly rather than running check_correctness.
    """
    x_shape = [1, 3, 4, 4]
    grid_shape = [1, 2, 2, 2]

    node = helper.make_node(
        "GridSample",
        inputs=["X", "grid"],
        outputs=["Y"],
        mode="cubic",
    )

    graph = helper.make_graph(
        [node],
        "grid_sample_cubic_test",
        inputs=[
            helper.make_tensor_value_info("X", TensorProto.FLOAT, x_shape),
            helper.make_tensor_value_info("grid", TensorProto.FLOAT, grid_shape),
        ],
        outputs=[
            helper.make_tensor_value_info(
                "Y", TensorProto.FLOAT, [x_shape[0], x_shape[1], grid_shape[1], grid_shape[2]]
            ),
        ],
    )

    model = helper.make_model(graph, producer_name="grid_sample_cubic_test")
    tvm_model = from_onnx(model, opset=16, keep_params_in_input=True)

    @I.ir_module
    class Expected:
        @R.function
        def main(
            X: R.Tensor((1, 3, 4, 4), dtype="float32"),
            grid: R.Tensor((1, 2, 2, 2), dtype="float32"),
        ) -> R.Tensor((1, 3, 2, 2), dtype="float32"):
            R.func_attr({"num_input": 2})
            with R.dataflow():
                lv: R.Tensor((1, 2, 2, 2), dtype="float32") = R.permute_dims(
                    grid, axes=[0, 3, 1, 2]
                )
                gv: R.Tensor((1, 3, 2, 2), dtype="float32") = R.image.grid_sample(
                    X,
                    lv,
                    method="bicubic",
                    layout="NCHW",
                    padding_mode="zeros",
                    align_corners=False,
                )
                R.output(gv)
            return gv

    tvm.ir.assert_structural_equal(tvm_model, Expected)


def test_roi_align():
    def verify_roi_align(coordinate_transformation_mode, rois, expected):
        x_shape = [1, 4, 8, 8]
        rois_shape = list(rois.shape)
        batch_indices_shape = [2]
        out_shape = [2, 4, 3, 3]

        node = helper.make_node(
            "RoiAlign",
            inputs=["X", "rois", "batch_indices"],
            outputs=["Y"],
            output_height=3,
            output_width=3,
            sampling_ratio=2,
            spatial_scale=1.0,
            mode="avg",
            coordinate_transformation_mode=coordinate_transformation_mode,
        )

        graph = helper.make_graph(
            [node],
            "roi_align_test",
            inputs=[
                helper.make_tensor_value_info("X", TensorProto.FLOAT, x_shape),
                helper.make_tensor_value_info("rois", TensorProto.FLOAT, rois_shape),
                helper.make_tensor_value_info(
                    "batch_indices", TensorProto.INT64, batch_indices_shape
                ),
            ],
            outputs=[helper.make_tensor_value_info("Y", TensorProto.FLOAT, out_shape)],
        )

        model = helper.make_model(graph, producer_name="roi_align_test")
        tvm_model = from_onnx(model, opset=16, keep_params_in_input=True)
        tvm.ir.assert_structural_equal(tvm_model, expected)

    @I.ir_module
    class ExpectedRoiAlignHalfPixel:
        @R.function
        def main(
            X: R.Tensor((1, 4, 8, 8), dtype="float32"),
            rois: R.Tensor((2, 4), dtype="float32"),
            batch_indices: R.Tensor((2,), dtype="int64"),
        ) -> R.Tensor((2, 4, 3, 3), dtype="float32"):
            R.func_attr({"num_input": 3})
            with R.dataflow():
                lv: R.Tensor((2, 1), dtype="int64") = R.expand_dims(batch_indices, axis=1)
                lv1: R.Tensor((2, 1), dtype="float32") = R.astype(lv, dtype="float32")
                lv2: R.Tensor((2, 4), dtype="float32") = R.add(
                    rois, R.const([-0.5, -0.5, -0.5, -0.5], "float32")
                )
                lv3: R.Tensor((2, 5), dtype="float32") = R.concat((lv1, lv2), axis=1)
                gv: R.Tensor((2, 4, 3, 3), dtype="float32") = R.vision.roi_align(
                    X,
                    lv3,
                    pooled_size=(3, 3),
                    spatial_scale=1.0,
                    sample_ratio=2,
                    aligned=True,
                    layout="NCHW",
                    mode="avg",
                )
                R.output(gv)
            return gv

    @I.ir_module
    class ExpectedRoiAlignOutputHalfPixel:
        @R.function
        def main(
            X: R.Tensor((1, 4, 8, 8), dtype="float32"),
            rois: R.Tensor((2, 4), dtype="float32"),
            batch_indices: R.Tensor((2,), dtype="int64"),
        ) -> R.Tensor((2, 4, 3, 3), dtype="float32"):
            R.func_attr({"num_input": 3})
            with R.dataflow():
                lv: R.Tensor((2, 1), dtype="int64") = R.expand_dims(batch_indices, axis=1)
                lv1: R.Tensor((2, 1), dtype="float32") = R.astype(lv, dtype="float32")
                lv2: R.Tensor((2, 5), dtype="float32") = R.concat((lv1, rois), axis=1)
                gv: R.Tensor((2, 4, 3, 3), dtype="float32") = R.vision.roi_align(
                    X,
                    lv2,
                    pooled_size=(3, 3),
                    spatial_scale=1.0,
                    sample_ratio=2,
                    aligned=False,
                    layout="NCHW",
                    mode="avg",
                )
                R.output(gv)
            return gv

    verify_roi_align(
        "output_half_pixel",
        np.array([[1.0, 1.0, 6.0, 6.0], [2.0, 0.5, 7.0, 7.0]], dtype="float32"),
        ExpectedRoiAlignOutputHalfPixel,
    )
    verify_roi_align(
        "half_pixel",
        np.array([[1.0, 1.0, 1.2, 1.2], [2.0, 0.5, 1.1, 1.1]], dtype="float32"),
        ExpectedRoiAlignHalfPixel,
    )


def test_if():
    """Test ONNX If operator with scalar and tensor bool conditions."""

    def verify_if(cond_info, expected):
        x_info = helper.make_tensor_value_info("x", TensorProto.FLOAT, [3])
        result_info = helper.make_tensor_value_info("result", TensorProto.FLOAT, [3])

        two = helper.make_tensor("two", TensorProto.FLOAT, [1], [2.0])
        then_mul = helper.make_node("Mul", ["x", "two"], ["then_out"])
        then_out_info = helper.make_tensor_value_info("then_out", TensorProto.FLOAT, [3])
        then_graph = helper.make_graph(
            [then_mul], "then_graph", [], [then_out_info], initializer=[two]
        )

        three = helper.make_tensor("three", TensorProto.FLOAT, [1], [3.0])
        else_mul = helper.make_node("Mul", ["x", "three"], ["else_out"])
        else_out_info = helper.make_tensor_value_info("else_out", TensorProto.FLOAT, [3])
        else_graph = helper.make_graph(
            [else_mul], "else_graph", [], [else_out_info], initializer=[three]
        )

        if_node = helper.make_node(
            "If",
            inputs=["cond"],
            outputs=["result"],
            then_branch=then_graph,
            else_branch=else_graph,
        )
        main_graph = helper.make_graph([if_node], "if_test", [cond_info, x_info], [result_info])
        model = helper.make_model(main_graph, opset_imports=[helper.make_opsetid("", 13)])
        tvm_model = from_onnx(model, keep_params_in_input=True)
        tvm.ir.assert_structural_equal(tvm_model, expected)

    @I.ir_module
    class ExpectedScalarCondition:
        @R.function
        def main(
            cond: R.Tensor((), dtype="bool"),
            x: R.Tensor((3,), dtype="float32"),
        ) -> R.Tensor((3,), dtype="float32"):
            R.func_attr({"num_input": 2})
            if cond:
                gv: R.Tensor((3,), dtype="float32") = R.multiply(x, R.const([2.0], "float32"))
                gv2: R.Tensor((3,), dtype="float32") = gv
            else:
                gv1: R.Tensor((3,), dtype="float32") = R.multiply(x, R.const([3.0], "float32"))
                gv2: R.Tensor((3,), dtype="float32") = gv1
            return gv2

    @I.ir_module
    class ExpectedTensorCondition:
        @R.function
        def main(
            cond: R.Tensor((1,), dtype="bool"),
            x: R.Tensor((3,), dtype="float32"),
        ) -> R.Tensor((3,), dtype="float32"):
            R.func_attr({"num_input": 2})
            if cond:
                gv: R.Tensor((3,), dtype="float32") = R.multiply(x, R.const([2.0], "float32"))
                gv2: R.Tensor((3,), dtype="float32") = gv
            else:
                gv1: R.Tensor((3,), dtype="float32") = R.multiply(x, R.const([3.0], "float32"))
                gv2: R.Tensor((3,), dtype="float32") = gv1
            return gv2

    verify_if(helper.make_tensor_value_info("cond", TensorProto.BOOL, []), ExpectedScalarCondition)
    verify_if(helper.make_tensor_value_info("cond", TensorProto.BOOL, [1]), ExpectedTensorCondition)


def test_if_computed_condition():
    """Test If where condition is computed from another op in the main graph."""
    x_info = helper.make_tensor_value_info("x", TensorProto.FLOAT, [3])
    result_info = helper.make_tensor_value_info("result", TensorProto.FLOAT, [3])

    zero = helper.make_tensor("zero", TensorProto.FLOAT, [], [0.0])
    reduce_node = helper.make_node(
        "ReduceSum", ["x"], ["x_sum"], keepdims=0, noop_with_empty_axes=0
    )
    greater_node = helper.make_node("Greater", ["x_sum", "zero"], ["cond"])

    two = helper.make_tensor("two", TensorProto.FLOAT, [1], [2.0])
    then_mul = helper.make_node("Mul", ["x", "two"], ["then_out"])
    then_out_info = helper.make_tensor_value_info("then_out", TensorProto.FLOAT, [3])
    then_graph = helper.make_graph([then_mul], "then_graph", [], [then_out_info], initializer=[two])

    three = helper.make_tensor("three", TensorProto.FLOAT, [1], [3.0])
    else_mul = helper.make_node("Mul", ["x", "three"], ["else_out"])
    else_out_info = helper.make_tensor_value_info("else_out", TensorProto.FLOAT, [3])
    else_graph = helper.make_graph(
        [else_mul], "else_graph", [], [else_out_info], initializer=[three]
    )

    if_node = helper.make_node(
        "If", inputs=["cond"], outputs=["result"], then_branch=then_graph, else_branch=else_graph
    )

    main_graph = helper.make_graph(
        [reduce_node, greater_node, if_node],
        "if_computed_cond",
        [x_info],
        [result_info],
        initializer=[zero],
    )
    model = helper.make_model(main_graph, opset_imports=[helper.make_opsetid("", 13)])

    tvm_model = from_onnx(model, keep_params_in_input=True)
    assert len(tvm_model["main"].attrs["params"]) == 1
    tvm_model["main"] = tvm_model["main"].without_attr("params")

    @I.ir_module
    class Expected:
        @R.function
        def main(
            x: R.Tensor((3,), dtype="float32"),
            zer: R.Tensor((), dtype="float32"),
        ) -> R.Tensor((3,), dtype="float32"):
            R.func_attr({"num_input": 1})
            gv: R.Tensor((), dtype="float32") = R.sum(x, axis=None, keepdims=False)
            gv1: R.Tensor((), dtype="bool") = R.greater(gv, zer)
            if gv1:
                gv2: R.Tensor((3,), dtype="float32") = R.multiply(x, R.const([2.0], "float32"))
                gv4: R.Tensor((3,), dtype="float32") = gv2
            else:
                gv3: R.Tensor((3,), dtype="float32") = R.multiply(x, R.const([3.0], "float32"))
                gv4: R.Tensor((3,), dtype="float32") = gv3
            return gv4

    tvm.ir.assert_structural_equal(tvm_model, Expected)


def test_if_multiple_outputs():
    """Test If operator where branches return multiple outputs."""
    cond_info = helper.make_tensor_value_info("cond", TensorProto.BOOL, [])
    x_info = helper.make_tensor_value_info("x", TensorProto.FLOAT, [3])
    out1_info = helper.make_tensor_value_info("out1", TensorProto.FLOAT, [3])
    out2_info = helper.make_tensor_value_info("out2", TensorProto.FLOAT, [3])

    two = helper.make_tensor("two", TensorProto.FLOAT, [1], [2.0])
    three = helper.make_tensor("three", TensorProto.FLOAT, [1], [3.0])

    then_mul1 = helper.make_node("Mul", ["x", "two"], ["then_out1"])
    then_mul2 = helper.make_node("Mul", ["x", "three"], ["then_out2"])
    then_o1 = helper.make_tensor_value_info("then_out1", TensorProto.FLOAT, [3])
    then_o2 = helper.make_tensor_value_info("then_out2", TensorProto.FLOAT, [3])
    then_graph = helper.make_graph(
        [then_mul1, then_mul2], "then_graph", [], [then_o1, then_o2], initializer=[two, three]
    )

    four = helper.make_tensor("four", TensorProto.FLOAT, [1], [4.0])
    five = helper.make_tensor("five", TensorProto.FLOAT, [1], [5.0])
    else_mul1 = helper.make_node("Mul", ["x", "four"], ["else_out1"])
    else_mul2 = helper.make_node("Mul", ["x", "five"], ["else_out2"])
    else_o1 = helper.make_tensor_value_info("else_out1", TensorProto.FLOAT, [3])
    else_o2 = helper.make_tensor_value_info("else_out2", TensorProto.FLOAT, [3])
    else_graph = helper.make_graph(
        [else_mul1, else_mul2], "else_graph", [], [else_o1, else_o2], initializer=[four, five]
    )

    if_node = helper.make_node(
        "If",
        inputs=["cond"],
        outputs=["out1", "out2"],
        then_branch=then_graph,
        else_branch=else_graph,
    )
    main_graph = helper.make_graph(
        [if_node], "if_multi_out", [cond_info, x_info], [out1_info, out2_info]
    )
    model = helper.make_model(main_graph, opset_imports=[helper.make_opsetid("", 13)])

    tvm_model = from_onnx(model, keep_params_in_input=True)

    @I.ir_module
    class Expected:
        @R.function
        def main(
            cond: R.Tensor((), dtype="bool"),
            x: R.Tensor((3,), dtype="float32"),
        ) -> R.Tuple(R.Tensor((3,), dtype="float32"), R.Tensor((3,), dtype="float32")):
            R.func_attr({"num_input": 2})
            if cond:
                gv: R.Tensor((3,), dtype="float32") = R.multiply(x, R.const([2.0], "float32"))
                gv1: R.Tensor((3,), dtype="float32") = R.multiply(x, R.const([3.0], "float32"))
                gv4: R.Tuple(
                    R.Tensor((3,), dtype="float32"),
                    R.Tensor((3,), dtype="float32"),
                ) = gv, gv1
            else:
                gv2: R.Tensor((3,), dtype="float32") = R.multiply(x, R.const([4.0], "float32"))
                gv3: R.Tensor((3,), dtype="float32") = R.multiply(x, R.const([5.0], "float32"))
                gv4: R.Tuple(
                    R.Tensor((3,), dtype="float32"),
                    R.Tensor((3,), dtype="float32"),
                ) = gv2, gv3
            gv5: R.Tensor((3,), dtype="float32") = gv4[0]
            gv6: R.Tensor((3,), dtype="float32") = gv4[1]
            return (gv5, gv6)

    tvm.ir.assert_structural_equal(tvm_model, Expected)


def test_if_nested():
    """Test nested If operator inside a branch."""
    cond1_info = helper.make_tensor_value_info("cond1", TensorProto.BOOL, [])
    cond2_info = helper.make_tensor_value_info("cond2", TensorProto.BOOL, [])
    x_info = helper.make_tensor_value_info("x", TensorProto.FLOAT, [3])
    result_info = helper.make_tensor_value_info("result", TensorProto.FLOAT, [3])

    # Inner then: x * 2
    two = helper.make_tensor("two", TensorProto.FLOAT, [1], [2.0])
    inner_then_mul = helper.make_node("Mul", ["x", "two"], ["inner_then_out"])
    inner_then_out_info = helper.make_tensor_value_info("inner_then_out", TensorProto.FLOAT, [3])
    inner_then_graph = helper.make_graph(
        [inner_then_mul], "inner_then", [], [inner_then_out_info], initializer=[two]
    )

    # Inner else: x * 3
    three = helper.make_tensor("three", TensorProto.FLOAT, [1], [3.0])
    inner_else_mul = helper.make_node("Mul", ["x", "three"], ["inner_else_out"])
    inner_else_out_info = helper.make_tensor_value_info("inner_else_out", TensorProto.FLOAT, [3])
    inner_else_graph = helper.make_graph(
        [inner_else_mul], "inner_else", [], [inner_else_out_info], initializer=[three]
    )

    # Outer then: nested If(cond2, x*2, x*3)
    inner_if = helper.make_node(
        "If",
        inputs=["cond2"],
        outputs=["outer_then_out"],
        then_branch=inner_then_graph,
        else_branch=inner_else_graph,
    )
    outer_then_out_info = helper.make_tensor_value_info("outer_then_out", TensorProto.FLOAT, [3])
    outer_then_graph = helper.make_graph([inner_if], "outer_then", [], [outer_then_out_info])

    # Outer else: x * 4
    four = helper.make_tensor("four", TensorProto.FLOAT, [1], [4.0])
    outer_else_mul = helper.make_node("Mul", ["x", "four"], ["outer_else_out"])
    outer_else_out_info = helper.make_tensor_value_info("outer_else_out", TensorProto.FLOAT, [3])
    outer_else_graph = helper.make_graph(
        [outer_else_mul], "outer_else", [], [outer_else_out_info], initializer=[four]
    )

    outer_if = helper.make_node(
        "If",
        inputs=["cond1"],
        outputs=["result"],
        then_branch=outer_then_graph,
        else_branch=outer_else_graph,
    )
    main_graph = helper.make_graph(
        [outer_if], "nested_if", [cond1_info, cond2_info, x_info], [result_info]
    )
    model = helper.make_model(main_graph, opset_imports=[helper.make_opsetid("", 13)])
    tvm_model = from_onnx(model, keep_params_in_input=True)

    @I.ir_module
    class Expected:
        @R.function
        def main(
            cond1: R.Tensor((), dtype="bool"),
            cond2: R.Tensor((), dtype="bool"),
            x: R.Tensor((3,), dtype="float32"),
        ) -> R.Tensor((3,), dtype="float32"):
            R.func_attr({"num_input": 3})
            if cond2:
                gv: R.Tensor((3,), dtype="float32") = R.multiply(x, R.const([2.0], "float32"))
                gv2: R.Tensor((3,), dtype="float32") = gv
            else:
                gv1: R.Tensor((3,), dtype="float32") = R.multiply(x, R.const([3.0], "float32"))
                gv2: R.Tensor((3,), dtype="float32") = gv1
            if cond1:
                gv4: R.Tensor((3,), dtype="float32") = gv2
            else:
                gv3: R.Tensor((3,), dtype="float32") = R.multiply(x, R.const([4.0], "float32"))
                gv4: R.Tensor((3,), dtype="float32") = gv3
            return gv4

    tvm.ir.assert_structural_equal(tvm_model, Expected)


# Helper that builds the ONNX graph for MatMulInteger so the tests don't repeat boilerplate code every time
def _make_matmulinteger_model(A_shape, B_shape, A_dtype, B_dtype, a_zp_array=None, b_zp_array=None):
    """Build a minimal single-node ONNX graph for MatMulInteger."""

    def np_dtype_to_onnx(dt):
        return {np.int8: TensorProto.INT8, np.uint8: TensorProto.UINT8}[dt]

    A_info = helper.make_tensor_value_info("A", np_dtype_to_onnx(A_dtype), A_shape)
    B_info = helper.make_tensor_value_info("B", np_dtype_to_onnx(B_dtype), B_shape)
    graph_inputs = [A_info, B_info]
    node_inputs = ["A", "B"]
    initializers = []

    def _add_zp(name, arr, dtype):
        onnx_dtype = np_dtype_to_onnx(dtype)
        shape = list(arr.shape)
        initializers.append(helper.make_tensor(name, onnx_dtype, shape, arr.flatten().tolist()))
        node_inputs.append(name)

    if a_zp_array is not None:
        _add_zp("a_zero_point", a_zp_array, A_dtype)
    elif b_zp_array is not None:
        node_inputs.append("")  # placeholder only needed if b_zp is present

    if b_zp_array is not None:
        _add_zp("b_zero_point", b_zp_array, B_dtype)

    out_info = helper.make_tensor_value_info("output", TensorProto.INT32, None)
    node = helper.make_node("MatMulInteger", inputs=node_inputs, outputs=["output"])
    graph = helper.make_graph(
        [node], "matmulinteger", graph_inputs, [out_info], initializer=initializers
    )
    model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 10)])
    model.ir_version = 8
    return model


def verify_matmulinteger_ir(A_shape, B_shape, A_dtype, B_dtype, expected, a_zp=None, b_zp=None):
    model = _make_matmulinteger_model(
        A_shape,
        B_shape,
        A_dtype,
        B_dtype,
        a_zp_array=np.array(a_zp, dtype=A_dtype) if a_zp is not None else None,
        b_zp_array=np.array(b_zp, dtype=B_dtype) if b_zp is not None else None,
    )
    tvm_model = from_onnx(model, opset=10, keep_params_in_input=True)
    if a_zp is not None or b_zp is not None:
        assert len(tvm_model["main"].attrs["params"]) == 2
        tvm_model["main"] = tvm_model["main"].without_attr("params")
    tvm.ir.assert_structural_equal(tvm_model, expected)


def test_matmulinteger():
    """2-D MatMulInteger should import dtype casts and zero-point subtraction."""

    @I.ir_module
    class ExpectedInt8:
        @R.function
        def main(
            A: R.Tensor((4, 8), dtype="int8"),
            B: R.Tensor((8, 6), dtype="int8"),
        ):
            R.func_attr({"num_input": 2})
            with R.dataflow():
                lv = R.astype(A, dtype="int32")
                lv1 = R.astype(B, dtype="int32")
                gv = R.matmul(lv, lv1, out_dtype="int32")
                R.output(gv)
            return gv

    @I.ir_module
    class ExpectedUInt8:
        @R.function
        def main(
            A: R.Tensor((4, 8), dtype="uint8"),
            B: R.Tensor((8, 6), dtype="uint8"),
        ):
            R.func_attr({"num_input": 2})
            with R.dataflow():
                lv = R.astype(A, dtype="int32")
                lv1 = R.astype(B, dtype="int32")
                gv = R.matmul(lv, lv1, out_dtype="int32")
                R.output(gv)
            return gv

    @I.ir_module
    class ExpectedUInt8Int8:
        @R.function
        def main(
            A: R.Tensor((4, 8), dtype="uint8"),
            B: R.Tensor((8, 6), dtype="int8"),
        ):
            R.func_attr({"num_input": 2})
            with R.dataflow():
                lv = R.astype(A, dtype="int32")
                lv1 = R.astype(B, dtype="int32")
                gv = R.matmul(lv, lv1, out_dtype="int32")
                R.output(gv)
            return gv

    @I.ir_module
    class ExpectedInt8UInt8:
        @R.function
        def main(
            A: R.Tensor((4, 8), dtype="int8"),
            B: R.Tensor((8, 6), dtype="uint8"),
        ):
            R.func_attr({"num_input": 2})
            with R.dataflow():
                lv = R.astype(A, dtype="int32")
                lv1 = R.astype(B, dtype="int32")
                gv = R.matmul(lv, lv1, out_dtype="int32")
                R.output(gv)
            return gv

    @I.ir_module
    class ExpectedUInt8ScalarZeroPoints:
        @R.function
        def main(
            A: R.Tensor((4, 8), dtype="uint8"),
            B: R.Tensor((8, 6), dtype="uint8"),
            a_zero_point: R.Tensor((), dtype="uint8"),
            b_zero_point: R.Tensor((), dtype="uint8"),
        ):
            R.func_attr({"num_input": 2})
            with R.dataflow():
                lv = R.astype(A, dtype="int32")
                lv1 = R.astype(a_zero_point, dtype="int32")
                lv2 = R.subtract(lv, lv1)
                lv3 = R.astype(B, dtype="int32")
                lv4 = R.astype(b_zero_point, dtype="int32")
                lv5 = R.subtract(lv3, lv4)
                gv = R.matmul(lv2, lv5, out_dtype="int32")
                R.output(gv)
            return gv

    @I.ir_module
    class ExpectedInt8ScalarZeroPoints:
        @R.function
        def main(
            A: R.Tensor((4, 8), dtype="int8"),
            B: R.Tensor((8, 6), dtype="int8"),
            a_zero_point: R.Tensor((), dtype="int8"),
            b_zero_point: R.Tensor((), dtype="int8"),
        ):
            R.func_attr({"num_input": 2})
            with R.dataflow():
                lv = R.astype(A, dtype="int32")
                lv1 = R.astype(a_zero_point, dtype="int32")
                lv2 = R.subtract(lv, lv1)
                lv3 = R.astype(B, dtype="int32")
                lv4 = R.astype(b_zero_point, dtype="int32")
                lv5 = R.subtract(lv3, lv4)
                gv = R.matmul(lv2, lv5, out_dtype="int32")
                R.output(gv)
            return gv

    verify_matmulinteger_ir([4, 8], [8, 6], np.int8, np.int8, ExpectedInt8)
    verify_matmulinteger_ir([4, 8], [8, 6], np.uint8, np.uint8, ExpectedUInt8)
    verify_matmulinteger_ir([4, 8], [8, 6], np.uint8, np.int8, ExpectedUInt8Int8)
    verify_matmulinteger_ir([4, 8], [8, 6], np.int8, np.uint8, ExpectedInt8UInt8)
    verify_matmulinteger_ir(
        [4, 8],
        [8, 6],
        np.uint8,
        np.uint8,
        ExpectedUInt8ScalarZeroPoints,
        a_zp=np.uint8(128),
        b_zp=np.uint8(128),
    )
    verify_matmulinteger_ir(
        [4, 8],
        [8, 6],
        np.int8,
        np.int8,
        ExpectedInt8ScalarZeroPoints,
        a_zp=np.int8(1),
        b_zp=np.int8(2),
    )


def test_matmulinteger_batched():
    """Batched MatMulInteger should import as batched Relax matmul."""

    @I.ir_module
    class ExpectedBatched3D:
        @R.function
        def main(
            A: R.Tensor((2, 4, 8), dtype="int8"),
            B: R.Tensor((2, 8, 6), dtype="int8"),
            a_zero_point: R.Tensor((), dtype="int8"),
            b_zero_point: R.Tensor((), dtype="int8"),
        ):
            R.func_attr({"num_input": 2})
            with R.dataflow():
                lv = R.astype(A, dtype="int32")
                lv1 = R.astype(a_zero_point, dtype="int32")
                lv2 = R.subtract(lv, lv1)
                lv3 = R.astype(B, dtype="int32")
                lv4 = R.astype(b_zero_point, dtype="int32")
                lv5 = R.subtract(lv3, lv4)
                gv = R.matmul(lv2, lv5, out_dtype="int32")
                R.output(gv)
            return gv

    @I.ir_module
    class ExpectedBatched4D:
        @R.function
        def main(
            A: R.Tensor((2, 3, 4, 8), dtype="int8"),
            B: R.Tensor((2, 3, 8, 6), dtype="int8"),
            a_zero_point: R.Tensor((), dtype="int8"),
            b_zero_point: R.Tensor((), dtype="int8"),
        ):
            R.func_attr({"num_input": 2})
            with R.dataflow():
                lv = R.astype(A, dtype="int32")
                lv1 = R.astype(a_zero_point, dtype="int32")
                lv2 = R.subtract(lv, lv1)
                lv3 = R.astype(B, dtype="int32")
                lv4 = R.astype(b_zero_point, dtype="int32")
                lv5 = R.subtract(lv3, lv4)
                gv = R.matmul(lv2, lv5, out_dtype="int32")
                R.output(gv)
            return gv

    verify_matmulinteger_ir(
        [2, 4, 8],
        [2, 8, 6],
        np.int8,
        np.int8,
        ExpectedBatched3D,
        a_zp=np.int8(1),
        b_zp=np.int8(2),
    )
    verify_matmulinteger_ir(
        [2, 3, 4, 8],
        [2, 3, 8, 6],
        np.int8,
        np.int8,
        ExpectedBatched4D,
        a_zp=np.int8(1),
        b_zp=np.int8(2),
    )


def test_matmulinteger_per_channel_zp():
    """1-D zero points should expand for per-row/per-column MatMulInteger."""

    @I.ir_module
    class ExpectedPerChannelZeroPoints:
        @R.function
        def main(
            A: R.Tensor((4, 8), dtype="int8"),
            B: R.Tensor((8, 6), dtype="int8"),
            a_zero_point: R.Tensor((4,), dtype="int8"),
            b_zero_point: R.Tensor((6,), dtype="int8"),
        ):
            R.func_attr({"num_input": 2})
            with R.dataflow():
                lv = R.astype(A, dtype="int32")
                lv1 = R.astype(a_zero_point, dtype="int32")
                lv2 = R.expand_dims(lv1, axis=-1)
                lv3 = R.subtract(lv, lv2)
                lv4 = R.astype(B, dtype="int32")
                lv5 = R.astype(b_zero_point, dtype="int32")
                lv6 = R.expand_dims(lv5, axis=0)
                lv7 = R.subtract(lv4, lv6)
                gv = R.matmul(lv3, lv7, out_dtype="int32")
                R.output(gv)
            return gv

    verify_matmulinteger_ir(
        [4, 8],
        [8, 6],
        np.int8,
        np.int8,
        ExpectedPerChannelZeroPoints,
        a_zp=np.arange(4, dtype=np.int8),
        b_zp=np.arange(6, dtype=np.int8),
    )


@pytest.mark.parametrize(
    ("pooled_shape", "rois"),
    [
        ((1, 1), np.array([[0.0, 1.0, 1.0, 6.0, 6.0], [0.0, 0.0, 0.0, 7.0, 7.0]], dtype="float32")),
        (
            (2, 3),
            np.array([[0.0, 1.2, 0.5, 6.8, 7.0], [0.0, -1.0, 2.0, 3.5, 5.2]], dtype="float32"),
        ),
        (
            (2, 2),
            np.array(
                [[0.0, 100.0, 100.0, 110.0, 110.0], [0.0, 1.0, 1.0, 6.0, 6.0]], dtype="float32"
            ),
        ),
    ],
)
def test_max_roi_pool(pooled_shape, rois):
    x_shape = [1, 4, 8, 8]
    out_shape = [2, 4, pooled_shape[0], pooled_shape[1]]

    node = helper.make_node(
        "MaxRoiPool",
        inputs=["X", "rois"],
        outputs=["Y"],
        pooled_shape=pooled_shape,
        spatial_scale=1.0,
    )

    graph = helper.make_graph(
        [node],
        "max_roi_pool_test",
        inputs=[
            helper.make_tensor_value_info("X", TensorProto.FLOAT, x_shape),
            helper.make_tensor_value_info("rois", TensorProto.FLOAT, [2, 5]),
        ],
        outputs=[helper.make_tensor_value_info("Y", TensorProto.FLOAT, out_shape)],
    )

    model = helper.make_model(graph, producer_name="max_roi_pool_test")
    inputs = {
        "X": rg.standard_normal(size=x_shape).astype("float32"),
        "rois": rois,
    }
    check_correctness(model, inputs=inputs, opset=16, rtol=1e-5, atol=1e-5)


def test_arg_min_max_select_last_index():
    """select_last_index=1 should lower to flip + argreduce + index remap."""

    def verify_select_last_index(op_name, axis, keepdims, expected):
        shape = [3, 4, 5]
        node = helper.make_node(
            op_name,
            inputs=["data"],
            outputs=["out"],
            axis=axis,
            keepdims=int(keepdims),
            select_last_index=1,
        )

        out_shape = list(shape)
        if keepdims:
            out_shape[axis] = 1
        else:
            out_shape.pop(axis)

        graph = helper.make_graph(
            [node],
            "arg_select_last_index_test",
            inputs=[helper.make_tensor_value_info("data", TensorProto.FLOAT, shape)],
            outputs=[helper.make_tensor_value_info("out", TensorProto.INT64, out_shape)],
        )
        model = helper.make_model(graph, producer_name="arg_select_last_index_test")
        tvm_model = from_onnx(model, opset=12, keep_params_in_input=True)
        tvm.ir.assert_structural_equal(tvm_model, expected)

    def make_expected(op_name, axis, keepdims):
        axis_extent = [3, 4, 5][axis] - 1
        reduce_op = R.argmax if op_name == "ArgMax" else R.argmin

        @I.ir_module
        class ExpectedArgReduceSelectLast:
            @R.function
            def main(
                data: R.Tensor((3, 4, 5), dtype="float32"),
            ):
                R.func_attr({"num_input": 1})
                with R.dataflow():
                    lv: R.Tensor((3, 4, 5), dtype="float32") = R.flip(data, axis=axis)
                    lv1 = reduce_op(lv, axis=axis, keepdims=keepdims)
                    gv = R.subtract(R.const(axis_extent, "int64"), lv1)
                    R.output(gv)
                return gv

        return ExpectedArgReduceSelectLast

    for op_name in ["ArgMax", "ArgMin"]:
        for axis in [0, 1, 2]:
            for keepdims in [True, False]:
                verify_select_last_index(
                    op_name, axis, keepdims, make_expected(op_name, axis, keepdims)
                )


def test_arg_min_max_select_last_index_no_tie():
    """select_last_index=0 should keep direct argreduce lowering."""

    def verify_no_tie(op_name, expected):
        shape = [4, 5]
        node = helper.make_node(
            op_name,
            inputs=["data"],
            outputs=["out"],
            axis=1,
            keepdims=1,
            select_last_index=0,
        )
        graph = helper.make_graph(
            [node],
            "arg_no_tie_test",
            inputs=[helper.make_tensor_value_info("data", TensorProto.FLOAT, shape)],
            outputs=[helper.make_tensor_value_info("out", TensorProto.INT64, [4, 1])],
        )
        model = helper.make_model(graph, producer_name="arg_no_tie_test")
        tvm_model = from_onnx(model, opset=12, keep_params_in_input=True)
        tvm.ir.assert_structural_equal(tvm_model, expected)

    @I.ir_module
    class ExpectedArgMax:
        @R.function
        def main(
            data: R.Tensor((4, 5), dtype="float32"),
        ) -> R.Tensor((4, 1), dtype="int64"):
            R.func_attr({"num_input": 1})
            with R.dataflow():
                gv: R.Tensor((4, 1), dtype="int64") = R.argmax(data, axis=1, keepdims=True)
                R.output(gv)
            return gv

    @I.ir_module
    class ExpectedArgMin:
        @R.function
        def main(
            data: R.Tensor((4, 5), dtype="float32"),
        ) -> R.Tensor((4, 1), dtype="int64"):
            R.func_attr({"num_input": 1})
            with R.dataflow():
                gv: R.Tensor((4, 1), dtype="int64") = R.argmin(data, axis=1, keepdims=True)
                R.output(gv)
            return gv

    verify_no_tie("ArgMax", ExpectedArgMax)
    verify_no_tie("ArgMin", ExpectedArgMin)


def test_arg_min_max_select_last_index_ir():
    """select_last_index=1 must lower to flip + argmax/argmin + subtract in the Relax IR."""

    def verify_select_last_index_ir(op_name, expected):
        shape = [3, 4, 5]
        node = helper.make_node(
            op_name,
            inputs=["data"],
            outputs=["out"],
            axis=1,
            keepdims=1,
            select_last_index=1,
        )
        graph = helper.make_graph(
            [node],
            "arg_select_last_index_ir_test",
            inputs=[helper.make_tensor_value_info("data", TensorProto.FLOAT, shape)],
            outputs=[helper.make_tensor_value_info("out", TensorProto.INT64, [3, 1, 5])],
        )
        model = helper.make_model(graph, producer_name="arg_select_last_index_ir_test")
        tvm_model = from_onnx(model, opset=12, keep_params_in_input=True)
        tvm.ir.assert_structural_equal(tvm_model, expected)

    @I.ir_module
    class ExpectedArgMax:
        @R.function
        def main(
            data: R.Tensor((3, 4, 5), dtype="float32"),
        ) -> R.Tensor((3, 1, 5), dtype="int64"):
            R.func_attr({"num_input": 1})
            with R.dataflow():
                lv: R.Tensor((3, 4, 5), dtype="float32") = R.flip(data, axis=1)
                lv1: R.Tensor((3, 1, 5), dtype="int64") = R.argmax(lv, axis=1, keepdims=True)
                gv: R.Tensor((3, 1, 5), dtype="int64") = R.subtract(R.const(3, "int64"), lv1)
                R.output(gv)
            return gv

    @I.ir_module
    class ExpectedArgMin:
        @R.function
        def main(
            data: R.Tensor((3, 4, 5), dtype="float32"),
        ) -> R.Tensor((3, 1, 5), dtype="int64"):
            R.func_attr({"num_input": 1})
            with R.dataflow():
                lv: R.Tensor((3, 4, 5), dtype="float32") = R.flip(data, axis=1)
                lv1: R.Tensor((3, 1, 5), dtype="int64") = R.argmin(lv, axis=1, keepdims=True)
                gv: R.Tensor((3, 1, 5), dtype="int64") = R.subtract(R.const(3, "int64"), lv1)
                R.output(gv)
            return gv

    verify_select_last_index_ir("ArgMax", ExpectedArgMax)
    verify_select_last_index_ir("ArgMin", ExpectedArgMin)


def test_split_to_sequence_keepdims_0():
    """keepdims=0, no split input: each chunk of size 1 has the split axis squeezed out."""

    def verify_split_to_sequence_keepdims_0(axis: int, expected):
        shape = [3, 4, 5]
        out_shape = [s for i, s in enumerate(shape) if i != axis]

        split_to_seq_node = helper.make_node(
            "SplitToSequence",
            ["data"],
            ["output"],
            axis=axis,
            keepdims=0,
        )
        graph = helper.make_graph(
            [split_to_seq_node],
            f"test_split_to_sequence_keepdims_0_axis{axis}",
            inputs=[helper.make_tensor_value_info("data", TensorProto.FLOAT, shape)],
            outputs=[
                helper.make_tensor_sequence_value_info("output", TensorProto.FLOAT, out_shape)
            ],
        )
        model = helper.make_model(graph, producer_name="test_split_to_sequence_keepdims_0")
        tvm_model = from_onnx(model, keep_params_in_input=True)
        tvm.ir.assert_structural_equal(tvm_model, expected)

    @I.ir_module
    class ExpectedKeepdims0Axis0:
        @R.function
        def main(
            data: R.Tensor((3, 4, 5), dtype="float32"),
        ) -> R.Tuple(
            R.Tensor((4, 5), dtype="float32"),
            R.Tensor((4, 5), dtype="float32"),
            R.Tensor((4, 5), dtype="float32"),
        ):
            R.func_attr({"num_input": 1})
            with R.dataflow():
                lv: R.Tuple(
                    R.Tensor((1, 4, 5), dtype="float32"),
                    R.Tensor((1, 4, 5), dtype="float32"),
                    R.Tensor((1, 4, 5), dtype="float32"),
                ) = R.split(data, indices_or_sections=3, axis=0)
                lv1: R.Tensor((1, 4, 5), dtype="float32") = lv[0]
                lv2: R.Tensor((1, 4, 5), dtype="float32") = lv[1]
                lv3: R.Tensor((1, 4, 5), dtype="float32") = lv[2]
                lv4: R.Tensor((4, 5), dtype="float32") = R.squeeze(lv1, axis=[0])
                lv5: R.Tensor((4, 5), dtype="float32") = R.squeeze(lv2, axis=[0])
                lv6: R.Tensor((4, 5), dtype="float32") = R.squeeze(lv3, axis=[0])
                gv: R.Tuple(
                    R.Tensor((4, 5), dtype="float32"),
                    R.Tensor((4, 5), dtype="float32"),
                    R.Tensor((4, 5), dtype="float32"),
                ) = lv4, lv5, lv6
                R.output(gv)
            return gv

    @I.ir_module
    class ExpectedKeepdims0Axis1:
        @R.function
        def main(
            data: R.Tensor((3, 4, 5), dtype="float32"),
        ) -> R.Tuple(
            R.Tensor((3, 5), dtype="float32"),
            R.Tensor((3, 5), dtype="float32"),
            R.Tensor((3, 5), dtype="float32"),
            R.Tensor((3, 5), dtype="float32"),
        ):
            R.func_attr({"num_input": 1})
            with R.dataflow():
                lv: R.Tuple(
                    R.Tensor((3, 1, 5), dtype="float32"),
                    R.Tensor((3, 1, 5), dtype="float32"),
                    R.Tensor((3, 1, 5), dtype="float32"),
                    R.Tensor((3, 1, 5), dtype="float32"),
                ) = R.split(data, indices_or_sections=4, axis=1)
                lv1: R.Tensor((3, 1, 5), dtype="float32") = lv[0]
                lv2: R.Tensor((3, 1, 5), dtype="float32") = lv[1]
                lv3: R.Tensor((3, 1, 5), dtype="float32") = lv[2]
                lv4: R.Tensor((3, 1, 5), dtype="float32") = lv[3]
                lv5: R.Tensor((3, 5), dtype="float32") = R.squeeze(lv1, axis=[1])
                lv6: R.Tensor((3, 5), dtype="float32") = R.squeeze(lv2, axis=[1])
                lv7: R.Tensor((3, 5), dtype="float32") = R.squeeze(lv3, axis=[1])
                lv8: R.Tensor((3, 5), dtype="float32") = R.squeeze(lv4, axis=[1])
                gv: R.Tuple(
                    R.Tensor((3, 5), dtype="float32"),
                    R.Tensor((3, 5), dtype="float32"),
                    R.Tensor((3, 5), dtype="float32"),
                    R.Tensor((3, 5), dtype="float32"),
                ) = lv5, lv6, lv7, lv8
                R.output(gv)
            return gv

    @I.ir_module
    class ExpectedKeepdims0Axis2:
        @R.function
        def main(
            data: R.Tensor((3, 4, 5), dtype="float32"),
        ) -> R.Tuple(
            R.Tensor((3, 4), dtype="float32"),
            R.Tensor((3, 4), dtype="float32"),
            R.Tensor((3, 4), dtype="float32"),
            R.Tensor((3, 4), dtype="float32"),
            R.Tensor((3, 4), dtype="float32"),
        ):
            R.func_attr({"num_input": 1})
            with R.dataflow():
                lv: R.Tuple(
                    R.Tensor((3, 4, 1), dtype="float32"),
                    R.Tensor((3, 4, 1), dtype="float32"),
                    R.Tensor((3, 4, 1), dtype="float32"),
                    R.Tensor((3, 4, 1), dtype="float32"),
                    R.Tensor((3, 4, 1), dtype="float32"),
                ) = R.split(data, indices_or_sections=5, axis=2)
                lv1: R.Tensor((3, 4, 1), dtype="float32") = lv[0]
                lv2: R.Tensor((3, 4, 1), dtype="float32") = lv[1]
                lv3: R.Tensor((3, 4, 1), dtype="float32") = lv[2]
                lv4: R.Tensor((3, 4, 1), dtype="float32") = lv[3]
                lv5: R.Tensor((3, 4, 1), dtype="float32") = lv[4]
                lv6: R.Tensor((3, 4), dtype="float32") = R.squeeze(lv1, axis=[2])
                lv7: R.Tensor((3, 4), dtype="float32") = R.squeeze(lv2, axis=[2])
                lv8: R.Tensor((3, 4), dtype="float32") = R.squeeze(lv3, axis=[2])
                lv9: R.Tensor((3, 4), dtype="float32") = R.squeeze(lv4, axis=[2])
                lv10: R.Tensor((3, 4), dtype="float32") = R.squeeze(lv5, axis=[2])
                gv: R.Tuple(
                    R.Tensor((3, 4), dtype="float32"),
                    R.Tensor((3, 4), dtype="float32"),
                    R.Tensor((3, 4), dtype="float32"),
                    R.Tensor((3, 4), dtype="float32"),
                    R.Tensor((3, 4), dtype="float32"),
                ) = lv6, lv7, lv8, lv9, lv10
                R.output(gv)
            return gv

    verify_split_to_sequence_keepdims_0(0, ExpectedKeepdims0Axis0)
    verify_split_to_sequence_keepdims_0(1, ExpectedKeepdims0Axis1)
    verify_split_to_sequence_keepdims_0(2, ExpectedKeepdims0Axis2)


def test_split_to_sequence_keepdims_ignored_when_split_provided():
    """Per spec: keepdims is ignored when split input is provided.
    TVM follows the spec — output keeps the split axis even with keepdims=0."""
    split_node = make_constant_node("split", TensorProto.INT64, (), [1])
    split_to_seq_node = helper.make_node(
        "SplitToSequence",
        ["data", "split"],
        ["output"],
        axis=0,
        keepdims=0,
    )
    graph = helper.make_graph(
        [split_node, split_to_seq_node],
        "test_split_to_sequence_keepdims_ignored",
        inputs=[helper.make_tensor_value_info("data", TensorProto.FLOAT, [4, 5])],
        outputs=[helper.make_tensor_sequence_value_info("output", TensorProto.FLOAT, [1, 5])],
    )
    model = helper.make_model(
        graph,
        producer_name="test_split_to_sequence_keepdims_ignored",
        opset_imports=[helper.make_opsetid("", 11)],
    )
    model.ir_version = 8
    tvm_model = from_onnx(model, opset=11, keep_params_in_input=True)

    @I.ir_module
    class Expected:
        @R.function
        def main(
            data: R.Tensor((4, 5), dtype="float32"),
        ) -> R.Tuple(
            R.Tensor((1, 5), dtype="float32"),
            R.Tensor((1, 5), dtype="float32"),
            R.Tensor((1, 5), dtype="float32"),
            R.Tensor((1, 5), dtype="float32"),
        ):
            R.func_attr({"num_input": 1})
            with R.dataflow():
                gv: R.Tuple(
                    R.Tensor((1, 5), dtype="float32"),
                    R.Tensor((1, 5), dtype="float32"),
                    R.Tensor((1, 5), dtype="float32"),
                    R.Tensor((1, 5), dtype="float32"),
                ) = R.split(data, indices_or_sections=4, axis=0)
                R.output(gv)
            return gv

    tvm.ir.assert_structural_equal(tvm_model, Expected)


def test_split_to_sequence_uneven_last_chunk():
    """Spec: last chunk may be smaller if dim is not divisible by scalar split."""

    def verify_split_to_sequence_uneven_last_chunk(axis: int, shape: list[int], expected):
        split_node = make_constant_node("split", TensorProto.INT64, (), [2])
        split_to_seq_node = helper.make_node(
            "SplitToSequence", ["data", "split"], ["output"], axis=axis, keepdims=1
        )
        graph = helper.make_graph(
            [split_node, split_to_seq_node],
            f"test_split_to_sequence_uneven_axis{axis}",
            inputs=[helper.make_tensor_value_info("data", TensorProto.FLOAT, shape)],
            outputs=[helper.make_tensor_sequence_value_info("output", TensorProto.FLOAT, None)],
        )
        model = helper.make_model(graph, producer_name="test_split_to_sequence_uneven")
        tvm_model = from_onnx(model, keep_params_in_input=True)
        tvm.ir.assert_structural_equal(tvm_model, expected)

    @I.ir_module
    class ExpectedUnevenAxis0:
        @R.function
        def main(
            data: R.Tensor((5, 4), dtype="float32"),
        ) -> R.Tuple(
            R.Tensor((2, 4), dtype="float32"),
            R.Tensor((2, 4), dtype="float32"),
            R.Tensor((1, 4), dtype="float32"),
        ):
            R.func_attr({"num_input": 1})
            with R.dataflow():
                gv: R.Tuple(
                    R.Tensor((2, 4), dtype="float32"),
                    R.Tensor((2, 4), dtype="float32"),
                    R.Tensor((1, 4), dtype="float32"),
                ) = R.split(data, indices_or_sections=3, axis=0)
                R.output(gv)
            return gv

    @I.ir_module
    class ExpectedUnevenAxis1:
        @R.function
        def main(
            data: R.Tensor((3, 5), dtype="float32"),
        ) -> R.Tuple(
            R.Tensor((3, 2), dtype="float32"),
            R.Tensor((3, 2), dtype="float32"),
            R.Tensor((3, 1), dtype="float32"),
        ):
            R.func_attr({"num_input": 1})
            with R.dataflow():
                gv: R.Tuple(
                    R.Tensor((3, 2), dtype="float32"),
                    R.Tensor((3, 2), dtype="float32"),
                    R.Tensor((3, 1), dtype="float32"),
                ) = R.split(data, indices_or_sections=3, axis=1)
                R.output(gv)
            return gv

    verify_split_to_sequence_uneven_last_chunk(0, [5, 4], ExpectedUnevenAxis0)
    verify_split_to_sequence_uneven_last_chunk(1, [3, 5], ExpectedUnevenAxis1)


def test_quantizelinear_singleton_qparams_opset10():
    """QuantizeLinear must treat shape-[1] scale/zp as scalar in opset10."""
    node = helper.make_node("QuantizeLinear", ["x", "scale", "zero_point"], ["y"])
    graph = helper.make_graph(
        [node],
        "quantizelinear_singleton_qparams_opset10",
        [helper.make_tensor_value_info("x", TensorProto.FLOAT, [4, 3, 2, 2])],
        [helper.make_tensor_value_info("y", TensorProto.UINT8, [4, 3, 2, 2])],
        initializer=[
            helper.make_tensor("scale", TensorProto.FLOAT, [1], [0.03125]),
            helper.make_tensor("zero_point", TensorProto.UINT8, [1], [127]),
        ],
    )
    model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 10)])

    x = rg.standard_normal((4, 3, 2, 2)).astype("float32")
    check_correctness(model, inputs={"x": x}, opset=10, check_dtypes=True)


def test_dequantizelinear_singleton_qparams_opset10():
    """DequantizeLinear must treat shape-[1] scale/zp as scalar in opset10."""
    node = helper.make_node("DequantizeLinear", ["x", "scale", "zero_point"], ["y"])
    graph = helper.make_graph(
        [node],
        "dequantizelinear_singleton_qparams_opset10",
        [helper.make_tensor_value_info("x", TensorProto.UINT8, [64])],
        [helper.make_tensor_value_info("y", TensorProto.FLOAT, [64])],
        initializer=[
            helper.make_tensor("scale", TensorProto.FLOAT, [1], [0.125]),
            helper.make_tensor("zero_point", TensorProto.UINT8, [1], [1]),
        ],
    )
    model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 10)])

    x = rg.integers(low=0, high=255, size=(64,), dtype=np.uint8)
    check_correctness(model, inputs={"x": x}, opset=10, check_dtypes=True)


def test_quantizelinear_optional_zero_point_opset13():
    """ONNX allows missing zero_point input; importer should default it to 0 (uint8)."""
    node = helper.make_node("QuantizeLinear", ["x", "scale"], ["y"])
    graph = helper.make_graph(
        [node],
        "quantizelinear_optional_zero_point_opset13",
        [helper.make_tensor_value_info("x", TensorProto.FLOAT, [2, 5])],
        [helper.make_tensor_value_info("y", TensorProto.UINT8, [2, 5])],
        initializer=[helper.make_tensor("scale", TensorProto.FLOAT, [], [0.2])],
    )
    model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 13)])

    x = rg.standard_normal((2, 5)).astype("float32")
    check_correctness(model, inputs={"x": x}, opset=13, check_dtypes=True)


def test_dynamicquantizelinear_opset11():
    """DynamicQuantizeLinear should import as quantization helper ops."""
    node = helper.make_node("DynamicQuantizeLinear", ["x"], ["y", "y_scale", "y_zero_point"])
    graph = helper.make_graph(
        [node],
        "dynamicquantizelinear_opset11",
        [helper.make_tensor_value_info("x", TensorProto.FLOAT, [2, 3, 4])],
        [
            helper.make_tensor_value_info("y", TensorProto.UINT8, [2, 3, 4]),
            helper.make_tensor_value_info("y_scale", TensorProto.FLOAT, []),
            helper.make_tensor_value_info("y_zero_point", TensorProto.UINT8, []),
        ],
    )
    model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 11)])

    tvm_model = from_onnx(model, opset=11, keep_params_in_input=True)

    @I.ir_module
    class Expected:
        @R.function
        def main(
            x: R.Tensor((2, 3, 4), dtype="float32"),
        ) -> R.Tuple(
            R.Tensor((2, 3, 4), dtype="uint8"),
            R.Tensor((), dtype="float32"),
            R.Tensor((), dtype="uint8"),
        ):
            R.func_attr({"num_input": 1})
            with R.dataflow():
                lv: R.Tensor((), dtype="float32") = R.max(x, axis=None, keepdims=False)
                lv1: R.Tensor((), dtype="float32") = R.maximum(R.const(0.0, "float32"), lv)
                lv2: R.Tensor((), dtype="float32") = R.min(x, axis=None, keepdims=False)
                lv3: R.Tensor((), dtype="float32") = R.minimum(R.const(0.0, "float32"), lv2)
                lv4: R.Tensor((), dtype="float32") = R.subtract(lv1, lv3)
                lv5: R.Tensor((), dtype="float32") = R.divide(lv4, R.const(255.0, "float32"))
                lv6: R.Tensor((), dtype="float32") = R.divide(lv3, lv5)
                lv7: R.Tensor((), dtype="float32") = R.subtract(R.const(0.0, "float32"), lv6)
                lv8: R.Tensor((), dtype="float32") = R.clip(lv7, R.prim_value(0), R.prim_value(255))
                lv9: R.Tensor((), dtype="float32") = R.round(lv8)
                lv10: R.Tensor((), dtype="uint8") = R.astype(lv9, dtype="uint8")
                lv11: R.Tensor((2, 3, 4), dtype="uint8") = R.quantize(
                    x, lv5, lv10, out_dtype="uint8", axis=0
                )
                gv: R.Tuple(
                    R.Tensor((2, 3, 4), dtype="uint8"),
                    R.Tensor((), dtype="float32"),
                    R.Tensor((), dtype="uint8"),
                ) = (lv11, lv5, lv10)
                R.output(gv)
            return gv

    tvm.ir.assert_structural_equal(tvm_model, Expected)


def test_quantizelinear_default_axis_opset10():
    """opset10 QuantizeLinear should honor default axis=1 (not hardcode axis=0)."""
    node = helper.make_node("QuantizeLinear", ["x", "scale", "zero_point"], ["y"])
    graph = helper.make_graph(
        [node],
        "quantizelinear_axis_opset10",
        [helper.make_tensor_value_info("x", TensorProto.FLOAT, [2, 3, 4])],
        [helper.make_tensor_value_info("y", TensorProto.UINT8, [2, 3, 4])],
        initializer=[
            helper.make_tensor("scale", TensorProto.FLOAT, [3], [0.05, 0.1, 0.2]),
            helper.make_tensor("zero_point", TensorProto.UINT8, [3], [1, 127, 250]),
        ],
    )
    model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 10)])

    x = rg.standard_normal((2, 3, 4)).astype("float32")
    check_correctness(model, inputs={"x": x}, opset=10, check_dtypes=True)


def test_dequantizelinear_default_axis_opset10():
    """opset10 DequantizeLinear should honor default axis=1 (not hardcode axis=0)."""
    node = helper.make_node("DequantizeLinear", ["x", "scale", "zero_point"], ["y"])
    graph = helper.make_graph(
        [node],
        "dequantizelinear_axis_opset10",
        [helper.make_tensor_value_info("x", TensorProto.UINT8, [2, 3, 4])],
        [helper.make_tensor_value_info("y", TensorProto.FLOAT, [2, 3, 4])],
        initializer=[
            helper.make_tensor("scale", TensorProto.FLOAT, [3], [0.05, 0.1, 0.2]),
            helper.make_tensor("zero_point", TensorProto.UINT8, [3], [1, 127, 250]),
        ],
    )
    model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 10)])

    x = rg.integers(low=0, high=255, size=(2, 3, 4), dtype=np.uint8)
    check_correctness(model, inputs={"x": x}, opset=10, check_dtypes=True)


if __name__ == "__main__":
    tvm.testing.main()