# 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.
# ruff: noqa: F841

import tvm.script
import tvm.testing
from tvm import relax
from tvm.ir import assert_structural_equal
from tvm.relax.testing.runtime_builtin import MakeShapeCode, MatchShapeCode
from tvm.script import ir as I
from tvm.script import relax as R
from tvm.script import tirx as T

# note: we expected RemovePurityChecking to be run first, so we force purity in most test cases


def test_const_shape_arg():
    MS = MatchShapeCode

    @tvm.script.ir_module
    class Before:
        @R.function
        def main(x: R.Shape([1, 2]), y: R.Shape):
            R.func_attr({"relax.force_pure": True})
            return x

        @T.prim_func(s_tir=True)
        def extra_func(H: T.Buffer(T.int64(4), "int64")):
            """Extra function, checks if the pass preserves it."""
            H[T.int64(1)] = H[T.int64(0)] + T.int64(1)

    @tvm.script.ir_module
    class Expected:
        @R.function
        def main(x: R.Shape([1, 2]), y: R.Shape):
            R.func_attr({"relax.force_pure": True})
            shape_heap = R.null_value()
            _ = R.call_packed("vm.builtin.check_shape_info", x, 2, "", ty_args=[R.Tuple()])
            _ = R.call_packed("vm.builtin.check_shape_info", y, -1, "", ty_args=[R.Tuple()])
            _ = R.call_packed(
                "vm.builtin.match_shape",
                x,
                shape_heap,
                2,
                MS.ASSERT_EQUAL_TO_IMM,
                1,
                MS.ASSERT_EQUAL_TO_IMM,
                2,
                "",
                ty_args=[R.Tuple()],
            )
            return x

        @T.prim_func(s_tir=True)
        def extra_func(H: T.Buffer(T.int64(4), "int64")):
            H[T.int64(1)] = H[T.int64(0)] + T.int64(1)

    before = Before
    expected = Expected
    after = relax.transform.VMShapeLower(emit_err_ctx=False)(before)
    assert_structural_equal(after, expected)


def test_static_fn_check():
    """Check static shape and function."""
    MS = MatchShapeCode

    @tvm.script.ir_module
    class Before:
        @R.function
        def main(f: R.Callable([R.Any], R.Any), y: R.Shape([1, 2])):
            R.func_attr({"relax.force_pure": True})
            return y

    @tvm.script.ir_module
    class Expected:
        @R.function
        def main(f: R.Callable([R.Any], R.Any), y: R.Shape([1, 2])):
            R.func_attr({"relax.force_pure": True})
            shape_heap = R.null_value()
            _ = R.call_packed("vm.builtin.check_func_info", f, "", ty_args=[R.Tuple()])
            _ = R.call_packed("vm.builtin.check_shape_info", y, 2, "", ty_args=[R.Tuple()])
            _ = R.call_packed(
                "vm.builtin.match_shape",
                y,
                shape_heap,
                2,
                MS.ASSERT_EQUAL_TO_IMM,
                1,
                MS.ASSERT_EQUAL_TO_IMM,
                2,
                "",
                ty_args=[R.Tuple()],
            )
            return y

    before = Before
    expected = Expected
    after = relax.transform.VMShapeLower(emit_err_ctx=False)(before)
    assert_structural_equal(after, expected)


def test_simple_symbolic_shape():
    MS = MatchShapeCode

    @tvm.script.ir_module
    class Before:
        @R.function
        def main(x: R.Tensor(["n", 2, "m"], "float32")):
            R.func_attr({"relax.force_pure": True})
            return x

    sindex = {
        "n": 0,
        "m": 1,
    }

    @tvm.script.ir_module
    class Expected:
        @R.function
        def main(x: R.Tensor(["n", 2, "m"], "float32")):
            R.func_attr({"relax.force_pure": True})
            shape_heap = R.call_builtin_with_ctx(
                "vm.builtin.alloc_shape_heap",
                [R.prim_value(2)],
                ty_args=[R.Tensor(ndim=1, dtype="int64")],
            )
            _ = R.call_packed(
                "vm.builtin.check_tensor_info",
                x,
                3,
                R.dtype("float32"),
                "",
                ty_args=[R.Tuple()],
            )
            _ = R.call_packed(
                "vm.builtin.match_shape",
                x,
                shape_heap,
                3,
                MS.STORE_TO_HEAP,
                sindex["n"],
                MS.ASSERT_EQUAL_TO_IMM,
                2,
                MS.STORE_TO_HEAP,
                sindex["m"],
                "",
                ty_args=[R.Tuple()],
            )
            return x

    before = Before
    expected = Expected
    after = relax.transform.VMShapeLower(emit_err_ctx=False)(before)
    assert_structural_equal(after, expected)


def test_symbolic_compute():
    MS = MatchShapeCode
    MK = MakeShapeCode

    @tvm.script.ir_module
    class Before:
        @R.function
        def main(x: R.Tensor(["n", "m"], "float32"), y: R.Tensor(ndim=3, dtype=None)) -> R.Shape(
            ndim=3
        ):
            R.func_attr({"relax.force_pure": True})
            m = T.int64()
            k = T.int64()
            z = R.match_cast(y, R.Tensor([k, m, k + 1], dtype=None))
            return R.shape([k + 1, m, 2])

    # slot assignment:
    # 0: n, 1: m, 2:k, 3: k+1
    sindex = {"n": 0, "m": 1, "k": 2, "k+1": 3}

    @tvm.script.ir_module
    class Expected:
        @T.prim_func(private=True, s_tir=True)
        def shape_func(H: T.Buffer(T.int64(4), "int64")):
            # generated compute function
            T.func_attr({"tirx.is_host_func": True})
            H[T.int64(sindex["k+1"])] = H[T.int64(sindex["k"])] + T.int64(1)

        @R.function
        def main(x: R.Tensor(["n", "m"], "float32"), y: R.Tensor(ndim=3, dtype=None)) -> R.Shape(
            ndim=3
        ):
            R.func_attr({"relax.force_pure": True})
            m = T.int64()
            k = T.int64()
            cls = Expected
            shape_heap = R.call_builtin_with_ctx(
                "vm.builtin.alloc_shape_heap",
                [R.prim_value(4)],
                ty_args=[R.Tensor(ndim=1, dtype="int64")],
            )
            _ = R.call_packed(
                "vm.builtin.check_tensor_info",
                x,
                2,
                R.dtype("float32"),
                "",
                ty_args=[R.Tuple()],
            )
            gv = R.null_value()
            _ = R.call_packed("vm.builtin.check_tensor_info", y, 3, gv, "", ty_args=[R.Tuple()])
            _ = R.call_packed(
                "vm.builtin.match_shape",
                x,
                shape_heap,
                2,
                MS.STORE_TO_HEAP,
                sindex["n"],
                MS.STORE_TO_HEAP,
                sindex["m"],
                "",
                ty_args=[R.Tuple()],
            )
            _ = R.call_packed(
                "vm.builtin.match_shape",
                y,
                shape_heap,
                3,
                MS.STORE_TO_HEAP,
                sindex["k"],
                MS.ASSERT_EQUAL_TO_LOAD,
                sindex["m"],
                MS.NO_OP,
                0,
                "",
                ty_args=[R.Tuple()],
            )
            _ = cls.shape_func(shape_heap)
            # extra assertion on y's shape after shape computation
            _ = R.call_packed(
                "vm.builtin.match_shape",
                y,
                shape_heap,
                3,
                MS.ASSERT_EQUAL_TO_LOAD,
                sindex["k"],
                MS.ASSERT_EQUAL_TO_LOAD,
                sindex["m"],
                MS.ASSERT_EQUAL_TO_LOAD,
                sindex["k+1"],
                "",
                ty_args=[R.Tuple()],
            )
            z = R.match_cast(y, R.Tensor([k, m, k + 1], dtype=None))
            # construct shape value for return
            s = R.call_packed(
                "vm.builtin.make_shape",
                shape_heap,
                3,
                MK.LOAD_SHAPE,
                sindex["k+1"],
                MK.LOAD_SHAPE,
                sindex["m"],
                MK.USE_IMM,
                2,
                ty_args=[R.Shape(ndim=3)],
            )
            return s

    before = Before
    expected = Expected
    after = relax.transform.VMShapeLower(emit_err_ctx=False)(before)
    assert_structural_equal(after, expected)


def test_tuple_handling():
    MS = MatchShapeCode

    @tvm.script.ir_module
    class Before:
        @R.function
        def main(
            x: R.Tuple(
                R.Tensor(["n", "m"], "float32"), R.Tuple(R.Shape, R.Tensor(["n", "k"], "int32"))
            ),
        ):
            R.func_attr({"relax.force_pure": True})
            return x

    # slot assignment:
    sindex = {"n": 0, "m": 1, "k": 2}

    @tvm.script.ir_module
    class Expected:
        @R.function
        def main(
            x: R.Tuple(
                R.Tensor(["n", "m"], "float32"), R.Tuple(R.Shape, R.Tensor(["n", "k"], "int32"))
            ),
        ):
            R.func_attr({"relax.force_pure": True})
            shape_heap = R.call_builtin_with_ctx(
                "vm.builtin.alloc_shape_heap",
                [R.prim_value(3)],
                ty_args=[R.Tensor(ndim=1, dtype="int64")],
            )
            # recursively unpack tuple for static info check
            _ = R.call_packed("vm.builtin.check_tuple_info", x, 2, "", ty_args=[R.Tuple()])
            t0 = x[0]
            _ = R.call_packed(
                "vm.builtin.check_tensor_info",
                t0,
                2,
                R.dtype("float32"),
                "",
                ty_args=[R.Tuple()],
            )
            t1 = x[1]
            _ = R.call_packed("vm.builtin.check_tuple_info", t1, 2, "", ty_args=[R.Tuple()])
            t1x0 = t1[0]
            _ = R.call_packed("vm.builtin.check_shape_info", t1x0, -1, "", ty_args=[R.Tuple()])
            t1x1 = t1[1]
            _ = R.call_packed(
                "vm.builtin.check_tensor_info",
                t1x1,
                2,
                R.dtype("int32"),
                "",
                ty_args=[R.Tuple()],
            )
            # match shape checks.
            _ = R.call_packed(
                "vm.builtin.match_shape",
                t0,
                shape_heap,
                2,
                MS.STORE_TO_HEAP,
                sindex["n"],
                MS.STORE_TO_HEAP,
                sindex["m"],
                "",
                ty_args=[R.Tuple()],
            )
            _ = R.call_packed(
                "vm.builtin.match_shape",
                t1x1,
                shape_heap,
                2,
                MS.ASSERT_EQUAL_TO_LOAD,
                sindex["n"],
                MS.STORE_TO_HEAP,
                sindex["k"],
                "",
                ty_args=[R.Tuple()],
            )
            return x

    before = Before
    expected = Expected
    after = relax.transform.VMShapeLower(emit_err_ctx=False)(before)
    assert_structural_equal(after, expected)


def test_return_match_check():
    """Test when return body is not same as ret_ty, runtime match check needed."""
    MS = MatchShapeCode

    @tvm.script.ir_module
    class Before:
        @R.function
        def main(x: R.Tensor(["n", "m"], "float32"), y: R.Any) -> R.Tuple(
            R.Tensor(["n", "m"], "float32")
        ):
            R.func_attr({"relax.force_pure": True})
            return y

    # slot assignment:
    sindex = {
        "n": 0,
        "m": 1,
    }

    @tvm.script.ir_module
    class Expected:
        @R.function
        def main(x: R.Tensor(["n", "m"], "float32"), y: R.Any) -> R.Tuple(
            R.Tensor(["n", "m"], "float32")
        ):
            R.func_attr({"relax.force_pure": True})
            shape_heap = R.call_builtin_with_ctx(
                "vm.builtin.alloc_shape_heap",
                [R.prim_value(2)],
                ty_args=[R.Tensor(ndim=1, dtype="int64")],
            )
            _ = R.call_packed(
                "vm.builtin.check_tensor_info", x, 2, R.dtype("float32"), "", ty_args=[R.Tuple()]
            )
            _ = R.call_packed(
                "vm.builtin.match_shape",
                x,
                shape_heap,
                2,
                MS.STORE_TO_HEAP,
                sindex["n"],
                MS.STORE_TO_HEAP,
                sindex["m"],
                "",
                ty_args=[R.Tuple()],
            )
            _ = R.call_packed("vm.builtin.check_tuple_info", y, 1, "", ty_args=[R.Tuple()])
            # emit runtime function call since y do not have the right type.
            y1 = R.call_packed("vm.builtin.tuple_getitem", y, 0, ty_args=[R.Any])
            # run check
            _ = R.call_packed(
                "vm.builtin.check_tensor_info",
                y1,
                2,
                R.dtype("float32"),
                "",
                ty_args=[R.Tuple()],
            )
            # shape check
            _ = R.call_packed(
                "vm.builtin.match_shape",
                y1,
                shape_heap,
                2,
                MS.ASSERT_EQUAL_TO_LOAD,
                sindex["n"],
                MS.ASSERT_EQUAL_TO_LOAD,
                sindex["m"],
                "",
                ty_args=[R.Tuple()],
            )

            return y

    before = Before
    expected = Expected
    after = relax.transform.VMShapeLower(emit_err_ctx=False)(before)
    assert_structural_equal(after, expected)


def test_return_match_check_with_new_expr():
    """Like test_return_match_check, but requires a computation

    When return body is not same as ret_ty, a runtime match
    check is required.  This match check may require a symbolic
    expression to be computed.
    """
    MS = MatchShapeCode

    @tvm.script.ir_module
    class Before:
        @R.function
        def main(x: R.Tensor(["n", "n"], "float32")) -> R.Tensor(["n * n"], "float32"):
            R.func_attr({"relax.force_pure": True})
            out = R.call_packed("flatten_matrix", x, ty_args=R.Any)
            return out

    # slot assignment:
    sindex = {
        "n": 0,
        "n * n": 1,
    }

    @tvm.script.ir_module
    class Expected:
        @R.function
        def main(x: R.Tensor(["n", "n"], "float32")) -> R.Tensor(["n * n"], "float32"):
            R.func_attr({"relax.force_pure": True})
            shape_heap = R.call_builtin_with_ctx(
                "vm.builtin.alloc_shape_heap",
                [R.prim_value(2)],
                ty_args=[R.Tensor(ndim=1, dtype="int64")],
            )
            _ = R.call_packed(
                "vm.builtin.check_tensor_info", x, 2, R.dtype("float32"), "", ty_args=[R.Tuple()]
            )
            _ = R.call_packed(
                "vm.builtin.match_shape",
                x,
                shape_heap,
                2,
                MS.STORE_TO_HEAP,
                sindex["n"],
                MS.ASSERT_EQUAL_TO_LOAD,
                sindex["n"],
                "",
                ty_args=[R.Tuple()],
            )

            _ = Expected.shape_func(shape_heap)

            out = R.call_packed("flatten_matrix", x, ty_args=R.Any)
            _ = R.call_packed(
                "vm.builtin.check_tensor_info",
                out,
                1,
                R.dtype("float32"),
                "",
                ty_args=[R.Tuple()],
            )
            _ = R.call_packed(
                "vm.builtin.match_shape",
                out,
                shape_heap,
                1,
                MS.ASSERT_EQUAL_TO_LOAD,
                sindex["n * n"],
                "",
                ty_args=[R.Tuple()],
            )
            return out

        @T.prim_func(private=True, s_tir=True)
        def shape_func(H: T.Buffer(T.int64(2), "int64")):
            # generated compute function
            T.func_attr({"tirx.is_host_func": True})
            H[T.int64(sindex["n * n"])] = H[T.int64(sindex["n"])] * H[T.int64(sindex["n"])]

    before = Before
    expected = Expected
    after = relax.transform.VMShapeLower(emit_err_ctx=False)(before)
    assert_structural_equal(after, expected)


def test_symbolic_shape_multiple_function():
    MS = MatchShapeCode
    MK = MakeShapeCode

    @I.ir_module
    class Before:
        @R.function
        def fn1(A: R.Tensor(("m", "n"), dtype="float32")):
            R.func_attr({"relax.force_pure": True})
            m = T.int64()
            n = T.int64()
            return A

        @R.function
        def fn2(A: R.Tensor(("n", "m"), dtype="float32")):
            R.func_attr({"relax.force_pure": True})
            n = T.int64()
            m = T.int64()
            return A

    # slot assignment:
    sindex_fn1 = {
        "m": 0,
        "n": 1,
    }
    sindex_fn2 = {
        "n": 0,
        "m": 1,
    }

    @I.ir_module
    class Expected:
        @R.function
        def fn1(A: R.Tensor(("m", "n"), dtype="float32")) -> R.Tensor(("m", "n"), dtype="float32"):
            R.func_attr({"relax.force_pure": True})
            m = T.int64()
            n = T.int64()
            shape_heap: R.Tensor(dtype="int64", ndim=1) = R.call_builtin_with_ctx(
                "vm.builtin.alloc_shape_heap",
                (R.prim_value(2),),
                ty_args=(R.Tensor(dtype="int64", ndim=1),),
            )
            _: R.Tuple = R.call_packed(
                "vm.builtin.check_tensor_info",
                A,
                R.prim_value(2),
                R.dtype("float32"),
                R.str(""),
                ty_args=(R.Tuple,),
            )
            _1: R.Tuple = R.call_packed(
                "vm.builtin.match_shape",
                A,
                shape_heap,
                R.prim_value(2),
                MS.STORE_TO_HEAP,
                sindex_fn1["m"],
                MS.STORE_TO_HEAP,
                sindex_fn1["n"],
                R.str(""),
                ty_args=(R.Tuple,),
            )
            return A

        @R.function
        def fn2(A: R.Tensor(("n", "m"), dtype="float32")) -> R.Tensor(("n", "m"), dtype="float32"):
            R.func_attr({"relax.force_pure": True})
            n = T.int64()
            m = T.int64()
            shape_heap: R.Tensor(dtype="int64", ndim=1) = R.call_builtin_with_ctx(
                "vm.builtin.alloc_shape_heap",
                (R.prim_value(2),),
                ty_args=(R.Tensor(dtype="int64", ndim=1),),
            )
            _2: R.Tuple = R.call_packed(
                "vm.builtin.check_tensor_info",
                A,
                R.prim_value(2),
                R.dtype("float32"),
                R.str(""),
                ty_args=(R.Tuple,),
            )
            _3: R.Tuple = R.call_packed(
                "vm.builtin.match_shape",
                A,
                shape_heap,
                R.prim_value(2),
                MS.STORE_TO_HEAP,
                sindex_fn2["n"],
                MS.STORE_TO_HEAP,
                sindex_fn2["m"],
                R.str(""),
                ty_args=(R.Tuple,),
            )
            return A

    before = Before
    expected = Expected
    after = relax.transform.VMShapeLower(emit_err_ctx=False)(before)
    assert_structural_equal(after, expected)


def test_check_lifted_weights():
    MS = MatchShapeCode

    @I.ir_module
    class Before:
        @R.function
        def main_transform_params(params: R.Tuple(R.Tensor((16, 16), dtype="float32"))) -> R.Tuple(
            R.Tensor((16, 16), dtype="float32")
        ):
            R.func_attr({"relax.force_pure": True})
            return params

        @R.function
        def main(x: R.Tensor((16, 16), "float32"), param_0: R.Tensor((16, 16), dtype="float32")):
            R.func_attr({"relax.force_pure": True, "num_input": 1})
            return (x, param_0)

    @I.ir_module
    class Expected:
        @R.function
        def main_transform_params(params: R.Tuple(R.Tensor((16, 16), dtype="float32"))) -> R.Tuple(
            R.Tensor((16, 16), dtype="float32")
        ):
            R.func_attr({"relax.force_pure": True})
            shape_heap: R.Any = R.null_value()
            _: R.Tuple = R.call_packed(
                "vm.builtin.check_tuple_info",
                params,
                R.prim_value(1),
                R.str(""),
                ty_args=(R.Tuple,),
            )
            gv: R.Tensor((16, 16), dtype="float32") = params[0]
            _1: R.Tuple = R.call_packed(
                "vm.builtin.check_tensor_info",
                gv,
                R.prim_value(2),
                R.dtype("float32"),
                R.str(""),
                ty_args=(R.Tuple,),
            )
            _2: R.Tuple = R.call_packed(
                "vm.builtin.match_shape",
                gv,
                shape_heap,
                R.prim_value(2),
                MS.ASSERT_EQUAL_TO_IMM,
                R.prim_value(16),
                MS.ASSERT_EQUAL_TO_IMM,
                R.prim_value(16),
                R.str(""),
                ty_args=(R.Tuple,),
            )
            return params

        @R.function
        def main(
            x: R.Tensor((16, 16), dtype="float32"), param_0: R.Tensor((16, 16), dtype="float32")
        ) -> R.Tuple(R.Tensor((16, 16), dtype="float32"), R.Tensor((16, 16), dtype="float32")):
            R.func_attr({"num_input": 1, "relax.force_pure": True})
            shape_heap: R.Any = R.null_value()
            _: R.Tuple = R.call_packed(
                "vm.builtin.check_tensor_info",
                x,
                R.prim_value(2),
                R.dtype("float32"),
                R.str(""),
                ty_args=(R.Tuple,),
            )
            _1: R.Tuple = R.call_packed(
                "vm.builtin.match_shape",
                x,
                shape_heap,
                R.prim_value(2),
                MS.ASSERT_EQUAL_TO_IMM,
                R.prim_value(16),
                MS.ASSERT_EQUAL_TO_IMM,
                R.prim_value(16),
                R.str(""),
                ty_args=(R.Tuple,),
            )
            return (x, param_0)

    before = Before
    after = relax.transform.VMShapeLower(emit_err_ctx=False)(before)
    expected = Expected
    assert_structural_equal(after, expected)


def test_check_weights_with_dynamic_shape():
    MS = MatchShapeCode

    @I.ir_module
    class Before:
        @R.function
        def main(
            x: R.Tensor((16, 16), "float32"),
            params: R.Tuple(R.Tensor((16, 16), dtype="float32"), R.Tensor(("n",), "float32")),
        ):
            R.func_attr({"relax.force_pure": True, "num_input": 1})
            n = T.int64()
            param_0 = params[0]
            param_1 = params[1]
            return (x, param_0, param_1)

    @I.ir_module
    class Expected:
        @R.function
        def main(
            x: R.Tensor((16, 16), "float32"),
            params: R.Tuple(R.Tensor((16, 16), dtype="float32"), R.Tensor(("n",), "float32")),
        ):
            n = T.int64()
            R.func_attr({"num_input": 1, "relax.force_pure": True})
            shape_heap: R.Tensor(dtype="int64", ndim=1) = R.call_builtin_with_ctx(
                "vm.builtin.alloc_shape_heap",
                (R.prim_value(1),),
                ty_args=(R.Tensor(dtype="int64", ndim=1),),
            )
            _: R.Tuple = R.call_packed(
                "vm.builtin.check_tensor_info",
                x,
                R.prim_value(2),
                R.dtype("float32"),
                R.str(""),
                ty_args=(R.Tuple,),
            )
            _1: R.Tuple = R.call_packed(
                "vm.builtin.check_tuple_info",
                params,
                R.prim_value(2),
                R.str(""),
                ty_args=(R.Tuple,),
            )
            _param_1: R.Tensor((n,), dtype="float32") = params[1]
            _2: R.Tuple = R.call_packed(
                "vm.builtin.check_tensor_info",
                _param_1,
                R.prim_value(1),
                R.dtype("float32"),
                R.str(""),
                ty_args=(R.Tuple,),
            )
            _3: R.Tuple = R.call_packed(
                "vm.builtin.match_shape",
                x,
                shape_heap,
                R.prim_value(2),
                MS.ASSERT_EQUAL_TO_IMM,
                R.prim_value(16),
                MS.ASSERT_EQUAL_TO_IMM,
                R.prim_value(16),
                R.str(""),
                ty_args=(R.Tuple,),
            )
            _4: R.Tuple = R.call_packed(
                "vm.builtin.match_shape",
                _param_1,
                shape_heap,
                MS.STORE_TO_HEAP,
                R.prim_value(1),
                R.prim_value(0),
                R.str(""),
                ty_args=(R.Tuple,),
            )

            param_0 = params[0]
            param_1 = params[1]
            return (x, param_0, param_1)

    before = Before
    after = relax.transform.VMShapeLower(emit_err_ctx=False)(before)
    print(after)
    expected = Expected
    assert_structural_equal(after, expected)


if __name__ == "__main__":
    tvm.testing.main()
