# 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 math
import operator

import pytest
import torch
import torch.nn.functional as F
from torch import fx
from torch.nn import Module

import tvm
import tvm.testing
from tvm import relax
from tvm.relax.frontend import detach_params
from tvm.relax.frontend.torch import from_fx
from tvm.script import ir as I
from tvm.script import relax as R
from tvm.script import tirx as T
from tvm.testing import env


def verify_model(torch_model, input_info, binding, expected):
    graph_model = fx.symbolic_trace(torch_model)
    with torch.no_grad():
        mod = from_fx(graph_model, input_info)
    binding = {k: tvm.runtime.tensor(v) for k, v in binding.items()}
    expected = relax.transform.BindParams("main", binding)(expected)
    tvm.ir.assert_structural_equal(mod, expected)


def test_conv1d():
    class Conv1D1(Module):
        def __init__(self):
            super().__init__()
            self.conv = torch.nn.Conv1d(3, 6, 7, bias=True)

        def forward(self, input):
            return self.conv(input)

    class Conv1D1Func(Module):
        def __init__(self):
            super().__init__()
            self.weight = torch.randn(size=[6, 3, 7])
            self.bias = torch.randn(size=[6])

        def forward(self, input):
            return torch.nn.functional.conv1d(input, self.weight, self.bias)

    @tvm.script.ir_module
    class expected1:
        @R.function
        def main(
            input_1: R.Tensor((1, 3, 10), dtype="float32"),
            w1: R.Tensor((6, 3, 7), dtype="float32"),
            w2: R.Tensor((6,), dtype="float32"),
        ) -> R.Tensor((1, 6, 4), dtype="float32"):
            # block 0
            with R.dataflow():
                lv1: R.Tensor((1, 6, 4), dtype="float32") = R.nn.conv1d(
                    input_1,
                    w1,
                    strides=[1],
                    padding=[0, 0],
                    dilation=[1],
                    data_layout="NCW",
                    kernel_layout="OIW",
                    out_layout="NCW",
                    out_dtype="float32",
                )
                lv2: R.Tensor((1, 6, 1), dtype="float32") = R.reshape(w2, [1, 6, 1])
                lv3: R.Tensor((1, 6, 4), dtype="float32") = R.add(lv1, lv2)
                gv: R.Tensor((1, 6, 4), dtype="float32") = lv3
                R.output(gv)
            return gv

    class Conv1D2(Module):
        def __init__(self):
            super().__init__()
            self.conv = torch.nn.Conv1d(3, 6, 7, bias=False)

        def forward(self, input):
            return self.conv(input)

    @tvm.script.ir_module
    class expected2:
        @R.function
        def main(
            input_1: R.Tensor((1, 3, 10), dtype="float32"),
            w1: R.Tensor((6, 3, 7), dtype="float32"),
        ) -> R.Tensor((1, 6, 4), dtype="float32"):
            # block 0
            with R.dataflow():
                lv1: R.Tensor((1, 6, 4), dtype="float32") = R.nn.conv1d(
                    input_1,
                    w1,
                    strides=[1],
                    padding=[0, 0],
                    dilation=[1],
                    data_layout="NCW",
                    kernel_layout="OIW",
                    out_layout="NCW",
                    out_dtype="float32",
                )
                gv: R.Tensor((1, 6, 4), dtype="float32") = lv1
                R.output(gv)
            return gv

    input_info = [([1, 3, 10], "float32")]

    model = Conv1D1()
    binding = {"w1": model.conv.weight.detach().numpy(), "w2": model.conv.bias.detach().numpy()}
    verify_model(model, input_info, binding, expected1)

    model = Conv1D1Func()
    binding = {"w1": model.weight.detach().numpy(), "w2": model.bias.detach().numpy()}
    verify_model(model, input_info, binding, expected1)

    model = Conv1D2()
    binding = {"w1": model.conv.weight.detach().numpy()}
    verify_model(model, input_info, binding, expected2)


def test_conv1d_transpose():
    class ConvTranspose1d1(Module):
        def __init__(self):
            super().__init__()
            self.conv = torch.nn.ConvTranspose1d(6, 6, 3, bias=True)

        def forward(self, input):
            return self.conv(input)

    class ConvTranspose1d1Func(Module):
        def __init__(self):
            super().__init__()
            self.weight = torch.randn(size=[6, 6, 3])
            self.bias = torch.randn(size=[6])

        def forward(self, input):
            return torch.nn.functional.conv_transpose1d(input, self.weight, self.bias)

    @tvm.script.ir_module
    class expected1:
        @R.function
        def main(
            input_1: R.Tensor((1, 6, 4), dtype="float32"),
            w1: R.Tensor((6, 6, 3), dtype="float32"),
            w2: R.Tensor((6,), dtype="float32"),
        ) -> R.Tensor((1, 6, 6), dtype="float32"):
            # block 0
            with R.dataflow():
                lv1: R.Tensor((1, 6, 6), dtype="float32") = R.nn.conv1d_transpose(
                    input_1,
                    w1,
                    strides=[1],
                    padding=[0, 0],
                    output_padding=[0],
                    dilation=[1],
                    data_layout="NCW",
                    kernel_layout="IOW",
                    out_layout="NCW",
                    out_dtype="float32",
                )
                lv2: R.Tensor((1, 6, 1), dtype="float32") = R.reshape(w2, [1, 6, 1])
                lv3: R.Tensor((1, 6, 6), dtype="float32") = R.add(lv1, lv2)
                gv: R.Tensor((1, 6, 6), dtype="float32") = lv3
                R.output(gv)
            return gv

    class ConvTranspose1d2(Module):
        def __init__(self):
            super().__init__()
            self.conv = torch.nn.ConvTranspose1d(6, 6, 3, bias=False)

        def forward(self, input):
            return self.conv(input)

    @tvm.script.ir_module
    class expected2:
        @R.function
        def main(
            input_1: R.Tensor((1, 6, 4), dtype="float32"),
            w1: R.Tensor((6, 6, 3), dtype="float32"),
        ) -> R.Tensor((1, 6, 6), dtype="float32"):
            # block 0
            with R.dataflow():
                lv1: R.Tensor((1, 6, 6), dtype="float32") = R.nn.conv1d_transpose(
                    input_1,
                    w1,
                    strides=[1],
                    padding=[0, 0],
                    output_padding=[0],
                    dilation=[1],
                    data_layout="NCW",
                    kernel_layout="IOW",
                    out_layout="NCW",
                    out_dtype="float32",
                )
                gv: R.Tensor((1, 6, 6), dtype="float32") = lv1
                R.output(gv)
            return gv

    input_info = [([1, 6, 4], "float32")]

    model = ConvTranspose1d1()
    binding = {"w1": model.conv.weight.detach().numpy(), "w2": model.conv.bias.detach().numpy()}
    verify_model(model, input_info, binding, expected1)

    model = ConvTranspose1d1Func()
    binding = {"w1": model.weight.detach().numpy(), "w2": model.bias.detach().numpy()}
    verify_model(model, input_info, binding, expected1)

    model = ConvTranspose1d2()
    binding = {"w1": model.conv.weight.detach().numpy()}
    verify_model(model, input_info, binding, expected2)


def test_conv2d():
    class Conv2D1(Module):
        def __init__(self):
            super().__init__()
            self.conv = torch.nn.Conv2d(3, 6, 7, bias=True)

        def forward(self, input):
            return self.conv(input)

    class Conv2D1Func(Module):
        def __init__(self):
            super().__init__()
            self.weight = torch.randn(size=[6, 3, 7, 7])
            self.bias = torch.randn(size=[6])

        def forward(self, input):
            return torch.nn.functional.conv2d(input, self.weight, self.bias)

    @tvm.script.ir_module
    class expected1:
        @R.function
        def main(
            input_1: R.Tensor((1, 3, 10, 10), dtype="float32"),
            w1: R.Tensor((6, 3, 7, 7), dtype="float32"),
            w2: R.Tensor((6,), dtype="float32"),
        ) -> R.Tensor((1, 6, 4, 4), dtype="float32"):
            # block 0
            with R.dataflow():
                lv1: R.Tensor((1, 6, 4, 4), dtype="float32") = R.nn.conv2d(
                    input_1,
                    w1,
                    strides=[1, 1],
                    padding=[0, 0, 0, 0],
                    dilation=[1, 1],
                    data_layout="NCHW",
                    kernel_layout="OIHW",
                    out_layout="NCHW",
                    out_dtype="float32",
                )
                lv2: R.Tensor((1, 6, 1, 1), dtype="float32") = R.reshape(w2, [1, 6, 1, 1])
                lv3: R.Tensor((1, 6, 4, 4), dtype="float32") = R.add(lv1, lv2)
                gv: R.Tensor((1, 6, 4, 4), dtype="float32") = lv3
                R.output(gv)
            return gv

    class Conv2D2(Module):
        def __init__(self):
            super().__init__()
            self.conv = torch.nn.Conv2d(3, 6, 7, bias=False)

        def forward(self, input):
            return self.conv(input)

    @tvm.script.ir_module
    class expected2:
        @R.function
        def main(
            input_1: R.Tensor((1, 3, 10, 10), dtype="float32"),
            w1: R.Tensor((6, 3, 7, 7), dtype="float32"),
        ) -> R.Tensor((1, 6, 4, 4), dtype="float32"):
            # block 0
            with R.dataflow():
                lv1: R.Tensor((1, 6, 4, 4), dtype="float32") = R.nn.conv2d(
                    input_1,
                    w1,
                    strides=[1, 1],
                    padding=[0, 0, 0, 0],
                    dilation=[1, 1],
                    data_layout="NCHW",
                    kernel_layout="OIHW",
                    out_layout="NCHW",
                    out_dtype="float32",
                )
                gv: R.Tensor((1, 6, 4, 4), dtype="float32") = lv1
                R.output(gv)
            return gv

    input_info = [([1, 3, 10, 10], "float32")]

    model = Conv2D1()
    binding = {"w1": model.conv.weight.detach().numpy(), "w2": model.conv.bias.detach().numpy()}
    verify_model(model, input_info, binding, expected1)

    model = Conv2D1Func()
    binding = {"w1": model.weight.numpy(), "w2": model.bias.numpy()}
    verify_model(model, input_info, binding, expected1)

    model = Conv2D2()
    binding = {"w1": model.conv.weight.detach().numpy()}
    verify_model(model, input_info, binding, expected2)


def test_conv2d_transpose():
    class ConvTranspose2d1(Module):
        def __init__(self):
            super().__init__()
            self.conv = torch.nn.ConvTranspose2d(3, 3, 7, bias=True)

        def forward(self, input):
            return self.conv(input)

    class ConvTranspose2d1Func(Module):
        def __init__(self):
            super().__init__()
            self.weight = torch.randn(size=[3, 3, 7, 7])
            self.bias = torch.randn(size=[3])

        def forward(self, input):
            return torch.nn.functional.conv_transpose2d(input, self.weight, self.bias)

    @tvm.script.ir_module
    class expected1:
        @R.function
        def main(
            input_1: R.Tensor((1, 3, 10, 10), dtype="float32"),
            w1: R.Tensor((3, 3, 7, 7), dtype="float32"),
            w2: R.Tensor((3,), dtype="float32"),
        ) -> R.Tensor((1, 3, 16, 16), dtype="float32"):
            # block 0
            with R.dataflow():
                lv1: R.Tensor((1, 3, 16, 16), dtype="float32") = R.nn.conv2d_transpose(
                    input_1,
                    w1,
                    strides=[1, 1],
                    padding=[0, 0, 0, 0],
                    output_padding=[0, 0],
                    dilation=[1, 1],
                    data_layout="NCHW",
                    kernel_layout="IOHW",
                    out_layout="NCHW",
                    out_dtype="float32",
                )
                lv2: R.Tensor((1, 3, 1, 1), dtype="float32") = R.reshape(w2, [1, 3, 1, 1])
                lv3: R.Tensor((1, 3, 16, 16), dtype="float32") = R.add(lv1, lv2)
                gv: R.Tensor((1, 3, 16, 16), dtype="float32") = lv3
                R.output(gv)
            return gv

    class ConvTranspose2d2(Module):
        def __init__(self):
            super().__init__()
            self.conv = torch.nn.ConvTranspose2d(3, 3, 7, bias=False)

        def forward(self, input):
            return self.conv(input)

    @tvm.script.ir_module
    class expected2:
        @R.function
        def main(
            input_1: R.Tensor((1, 3, 10, 10), dtype="float32"),
            w1: R.Tensor((3, 3, 7, 7), dtype="float32"),
        ) -> R.Tensor((1, 3, 16, 16), dtype="float32"):
            # block 0
            with R.dataflow():
                lv1: R.Tensor((1, 3, 16, 16), dtype="float32") = R.nn.conv2d_transpose(
                    input_1,
                    w1,
                    strides=[1, 1],
                    padding=[0, 0, 0, 0],
                    output_padding=[0, 0],
                    dilation=[1, 1],
                    data_layout="NCHW",
                    kernel_layout="IOHW",
                    out_layout="NCHW",
                    out_dtype="float32",
                )
                gv: R.Tensor((1, 3, 16, 16), dtype="float32") = lv1
                R.output(gv)
            return gv

    input_info = [([1, 3, 10, 10], "float32")]

    model = ConvTranspose2d1()
    binding = {"w1": model.conv.weight.detach().numpy(), "w2": model.conv.bias.detach().numpy()}
    verify_model(model, input_info, binding, expected1)

    model = ConvTranspose2d1Func()
    binding = {"w1": model.weight.detach().numpy(), "w2": model.bias.detach().numpy()}
    verify_model(model, input_info, binding, expected1)

    model = ConvTranspose2d2()
    binding = {"w1": model.conv.weight.detach().numpy()}
    verify_model(model, input_info, binding, expected2)


def test_conv3d():
    class Conv3D1(Module):
        def __init__(self):
            super().__init__()
            self.conv = torch.nn.Conv3d(3, 6, 7, bias=True)

        def forward(self, input):
            return self.conv(input)

    class Conv3D1Func(Module):
        def __init__(self):
            super().__init__()
            self.weight = torch.randn(size=[6, 3, 7, 7, 7])
            self.bias = torch.randn(size=[6])

        def forward(self, input):
            return torch.nn.functional.conv3d(input, self.weight, self.bias)

    @tvm.script.ir_module
    class expected1:
        @R.function
        def main(
            input_1: R.Tensor((1, 3, 10, 10, 10), dtype="float32"),
            w1: R.Tensor((6, 3, 7, 7, 7), dtype="float32"),
            w2: R.Tensor((6,), dtype="float32"),
        ) -> R.Tensor((1, 6, 4, 4, 4), dtype="float32"):
            # block 0
            with R.dataflow():
                lv1: R.Tensor((1, 6, 4, 4, 4), dtype="float32") = R.nn.conv3d(
                    input_1,
                    w1,
                    strides=[1],
                    padding=[0, 0, 0],
                    dilation=[1],
                    data_layout="NCDHW",
                    kernel_layout="OIDHW",
                    out_layout="NCDHW",
                    out_dtype="float32",
                )
                lv2: R.Tensor((1, 6, 1, 1, 1), dtype="float32") = R.reshape(w2, [1, 6, 1, 1, 1])
                lv3: R.Tensor((1, 6, 4, 4, 4), dtype="float32") = R.add(lv1, lv2)
                gv: R.Tensor((1, 6, 4, 4, 4), dtype="float32") = lv3
                R.output(gv)
            return gv

    class Conv3D2(Module):
        def __init__(self):
            super().__init__()
            self.conv = torch.nn.Conv3d(3, 6, 7, bias=False)

        def forward(self, input):
            return self.conv(input)

    @tvm.script.ir_module
    class expected2:
        @R.function
        def main(
            input_1: R.Tensor((1, 3, 10, 10, 10), dtype="float32"),
            w1: R.Tensor((6, 3, 7, 7, 7), dtype="float32"),
        ) -> R.Tensor((1, 6, 4, 4, 4), dtype="float32"):
            # block 0
            with R.dataflow():
                lv1: R.Tensor((1, 6, 4, 4, 4), dtype="float32") = R.nn.conv3d(
                    input_1,
                    w1,
                    strides=[1],
                    padding=[0, 0, 0],
                    dilation=[1],
                    data_layout="NCDHW",
                    kernel_layout="OIDHW",
                    out_layout="NCDHW",
                    out_dtype="float32",
                )
                gv: R.Tensor((1, 6, 4, 4, 4), dtype="float32") = lv1
                R.output(gv)
            return gv

    input_info = [([1, 3, 10, 10, 10], "float32")]

    model = Conv3D1()
    binding = {"w1": model.conv.weight.detach().numpy(), "w2": model.conv.bias.detach().numpy()}
    verify_model(model, input_info, binding, expected1)

    model = Conv3D1Func()
    binding = {"w1": model.weight.detach().numpy(), "w2": model.bias.detach().numpy()}
    verify_model(model, input_info, binding, expected1)

    model = Conv3D2()
    binding = {"w1": model.conv.weight.detach().numpy()}
    verify_model(model, input_info, binding, expected2)


def test_pad():
    class PadModel(torch.nn.Module):
        def __init__(self, pad, mode="constant", value=0.0):
            super().__init__()
            self.pad = pad
            self.mode = mode
            self.value = value

        def forward(self, x):
            if self.mode == "constant":
                return torch.nn.functional.pad(x, self.pad, mode=self.mode, value=self.value)
            else:
                return torch.nn.functional.pad(x, self.pad, mode=self.mode)

    @tvm.script.ir_module
    class expected_constant:
        @R.function
        def main(x: R.Tensor((1, 3, 10, 10), dtype="float32")) -> R.Tensor(
            (1, 3, 14, 12), dtype="float32"
        ):
            with R.dataflow():
                lv: R.Tensor((1, 3, 14, 12), dtype="float32") = R.nn.pad(
                    x,
                    pad_width=[0, 0, 0, 0, 2, 2, 1, 1],
                    pad_mode="constant",
                    pad_value=0.0,
                )
                gv: R.Tensor((1, 3, 14, 12), dtype="float32") = lv
                R.output(gv)
            return gv

    @tvm.script.ir_module
    class expected_reflect:
        @R.function
        def main(x: R.Tensor((1, 3, 10, 10), dtype="float32")) -> R.Tensor(
            (1, 3, 14, 12), dtype="float32"
        ):
            with R.dataflow():
                lv: R.Tensor((1, 3, 14, 12), dtype="float32") = R.nn.pad(
                    x,
                    pad_width=[0, 0, 0, 0, 2, 2, 1, 1],
                    pad_mode="reflect",
                    pad_value=0.0,
                )
                gv: R.Tensor((1, 3, 14, 12), dtype="float32") = lv
                R.output(gv)
            return gv

    @tvm.script.ir_module
    class expected_replicate:
        @R.function
        def main(x: R.Tensor((1, 3, 10, 10), dtype="float32")) -> R.Tensor(
            (1, 3, 14, 12), dtype="float32"
        ):
            with R.dataflow():
                lv: R.Tensor((1, 3, 14, 12), dtype="float32") = R.nn.pad(
                    x,
                    pad_width=[0, 0, 0, 0, 2, 2, 1, 1],
                    pad_mode="replicate",
                    pad_value=0.0,
                )
                gv: R.Tensor((1, 3, 14, 12), dtype="float32") = lv
                R.output(gv)
            return gv

    @tvm.script.ir_module
    class expected_circular:
        @R.function
        def main(x: R.Tensor((1, 3, 10, 10), dtype="float32")) -> R.Tensor(
            (1, 3, 14, 12), dtype="float32"
        ):
            with R.dataflow():
                lv: R.Tensor((1, 3, 14, 12), dtype="float32") = R.nn.pad(
                    x,
                    pad_width=[0, 0, 0, 0, 2, 2, 1, 1],
                    pad_mode="circular",
                    pad_value=0.0,
                )
                gv: R.Tensor((1, 3, 14, 12), dtype="float32") = lv
                R.output(gv)
            return gv

    input_infos = [([1, 3, 10, 10], "float32")]
    verify_model(PadModel(pad=[1, 1, 2, 2]), input_infos, {}, expected_constant)
    verify_model(PadModel(pad=[1, 1, 2, 2], mode="reflect"), input_infos, {}, expected_reflect)
    verify_model(PadModel(pad=[1, 1, 2, 2], mode="replicate"), input_infos, {}, expected_replicate)
    verify_model(PadModel(pad=[1, 1, 2, 2], mode="circular"), input_infos, {}, expected_circular)


def test_pixel_shuffle():
    class PixelShuffle1(torch.nn.Module):
        def __init__(self, upscale_factor=2):
            super().__init__()
            self.pixel_shuffle = torch.nn.PixelShuffle(upscale_factor)

        def forward(self, x):
            return self.pixel_shuffle(x)

    class PixelShuffle2(torch.nn.Module):
        def __init__(self, upscale_factor=2):
            super().__init__()
            self.upscale_factor = upscale_factor

        def forward(self, x):
            return torch.nn.functional.pixel_shuffle(x, self.upscale_factor)

    @tvm.script.ir_module
    class expected:
        @R.function
        def main(inp_0: R.Tensor((1, 8, 10, 15), dtype="float32")) -> R.Tensor(
            (1, 2, 20, 30), dtype="float32"
        ):
            with R.dataflow():
                lv: R.Tensor((1, 2, 20, 30), dtype="float32") = R.nn.pixel_shuffle(
                    inp_0, upscale_factor=2
                )
                gv: R.Tensor((1, 2, 20, 30), dtype="float32") = lv
                R.output(gv)
            return gv

    input_infos = [([1, 8, 10, 15], "float32")]
    verify_model(PixelShuffle1(2), input_infos, {}, expected)
    verify_model(PixelShuffle2(2), input_infos, {}, expected)


def test_linear():
    # nn.Linear
    class Dense1(Module):
        def __init__(self):
            super().__init__()
            self.linear = torch.nn.Linear(10, 7, bias=True)

        def forward(self, input):
            return self.linear(input)

    class Dense1Func(Module):
        def __init__(self):
            super().__init__()
            self.weight = torch.randn(size=[7, 10])
            self.bias = torch.randn(size=[7])

        def forward(self, input):
            return torch.nn.functional.linear(input, self.weight, self.bias)

    @tvm.script.ir_module
    class expected1:
        @R.function
        def main(
            input_1: R.Tensor((1, 3, 10, 10), dtype="float32"),
            w1: R.Tensor((7, 10), dtype="float32"),
            w2: R.Tensor((7,), dtype="float32"),
        ) -> R.Tensor((1, 3, 10, 7), dtype="float32"):
            # block 0
            with R.dataflow():
                lv: R.Tensor((10, 7), dtype="float32") = R.permute_dims(w1, axes=None)
                lv1: R.Tensor((1, 3, 10, 7), dtype="float32") = R.matmul(
                    input_1, lv, out_dtype="float32"
                )
                lv2: R.Tensor((1, 3, 10, 7), dtype="float32") = R.add(lv1, w2)
                gv: R.Tensor((1, 3, 10, 7), dtype="float32") = lv2
                R.output(gv)
            return gv

    class Dense2(Module):
        def __init__(self):
            super().__init__()
            self.linear = torch.nn.Linear(10, 7, bias=False)

        def forward(self, input):
            return self.linear(input)

    @tvm.script.ir_module
    class expected2:
        @R.function
        def main(
            input_1: R.Tensor((1, 3, 10, 10), dtype="float32"),
            w1: R.Tensor((7, 10), dtype="float32"),
        ) -> R.Tensor((1, 3, 10, 7), dtype="float32"):
            # block 0
            with R.dataflow():
                lv: R.Tensor((10, 7), dtype="float32") = R.permute_dims(w1, axes=None)
                lv1: R.Tensor((1, 3, 10, 7), dtype="float32") = R.matmul(
                    input_1, lv, out_dtype="float32"
                )
                gv: R.Tensor((1, 3, 10, 7), dtype="float32") = lv1
                R.output(gv)
            return gv

    input_info = [([1, 3, 10, 10], "float32")]

    model = Dense1()
    binding = {"w1": model.linear.weight.detach().numpy(), "w2": model.linear.bias.detach().numpy()}
    verify_model(model, input_info, binding, expected1)

    model = Dense1Func()
    binding = {"w1": model.weight.numpy(), "w2": model.bias.numpy()}
    verify_model(model, input_info, binding, expected1)

    model = Dense2()
    binding = {"w1": model.linear.weight.detach().numpy()}
    verify_model(model, input_info, binding, expected2)

    # matmul
    class MatMul1(Module):
        def __init__(self):
            super().__init__()

        def forward(self, x, y):
            return torch.matmul(x, y)

    @tvm.script.ir_module
    class expected3:
        @R.function
        def main(
            input_1: R.Tensor((10, 10), dtype="float32"),
            input_2: R.Tensor((10, 10), dtype="float32"),
        ) -> R.Tensor((10, 10), dtype="float32"):
            # block 0
            with R.dataflow():
                lv: R.Tensor((10, 10), dtype="float32") = R.matmul(
                    input_1, input_2, out_dtype="float32"
                )
                gv: R.Tensor((10, 10), dtype="float32") = lv
                R.output(gv)
            return gv

    verify_model(
        MatMul1(),
        [([10, 10], "float32"), ([10, 10], "float32")],
        {},
        expected3,
    )


def test_bmm():
    class BMM(Module):
        def __init__(self):
            super().__init__()

        def forward(self, x, y):
            return torch.bmm(x, y)

    @tvm.script.ir_module
    class Expected:
        @R.function
        def main(
            input_1: R.Tensor((4, 128, 256), dtype="float32"),
            input_2: R.Tensor((4, 256, 512), dtype="float32"),
        ) -> R.Tensor((4, 128, 512), dtype="float32"):
            # block 0
            with R.dataflow():
                lv: R.Tensor((4, 128, 512), dtype="float32") = R.matmul(
                    input_1, input_2, out_dtype="float32"
                )
                gv: R.Tensor((4, 128, 512), dtype="float32") = lv
                R.output(gv)
            return gv

    verify_model(
        BMM(),
        [((4, 128, 256), "float32"), ((4, 256, 512), "float32")],
        {},
        Expected,
    )


def test_baddbmm():
    class BAddBMM1(Module):
        def __init__(self):
            super().__init__()

        def forward(self, c, x, y):
            return torch.baddbmm(c, x, y)

    class BAddBMM2(Module):
        def __init__(self):
            super().__init__()

        def forward(self, c, x, y):
            return torch.baddbmm(c, x, y, alpha=2, beta=0)

    @tvm.script.ir_module
    class Expected1:
        @R.function
        def main(
            inp_0: R.Tensor((4, 128, 512), dtype="float32"),
            inp_1: R.Tensor((4, 128, 256), dtype="float32"),
            inp_2: R.Tensor((4, 256, 512), dtype="float32"),
        ) -> R.Tensor((4, 128, 512), dtype="float32"):
            with R.dataflow():
                lv: R.Tensor((4, 128, 512), dtype="float32") = R.matmul(inp_1, inp_2)
                lv1: R.Tensor((4, 128, 512), dtype="float32") = R.add(lv, inp_0)
                gv: R.Tensor((4, 128, 512), dtype="float32") = lv1
                R.output(gv)
            return gv

    @tvm.script.ir_module
    class Expected2:
        @R.function
        def main(
            inp_0: R.Tensor((4, 128, 512), dtype="float32"),
            inp_1: R.Tensor((4, 128, 256), dtype="float32"),
            inp_2: R.Tensor((4, 256, 512), dtype="float32"),
        ) -> R.Tensor((4, 128, 512), dtype="float32"):
            with R.dataflow():
                lv: R.Tensor((4, 128, 512), dtype="float32") = R.matmul(inp_1, inp_2)
                lv1: R.Tensor((4, 128, 512), dtype="float32") = R.multiply(
                    lv, R.const(2, "float32")
                )
                gv: R.Tensor((4, 128, 512), dtype="float32") = lv1
                R.output(gv)
            return gv

    verify_model(
        BAddBMM1(),
        [((4, 128, 512), "float32"), ((4, 128, 256), "float32"), ((4, 256, 512), "float32")],
        {},
        Expected1,
    )

    verify_model(
        BAddBMM2(),
        [((4, 128, 512), "float32"), ((4, 128, 256), "float32"), ((4, 256, 512), "float32")],
        {},
        Expected2,
    )


def test_einsum():
    class Einsum1(Module):
        def __init__(self):
            super().__init__()

        def forward(self, x):
            return torch.einsum("ii", x)

    class Einsum2(Module):
        def __init__(self):
            super().__init__()

        def forward(self, x, y):
            return torch.einsum("i,j->ij", x, y)

    @tvm.script.ir_module
    class Expected1:
        @R.function
        def main(inp_0: R.Tensor((4, 4), dtype="float32")) -> R.Tensor((), dtype="float32"):
            with R.dataflow():
                lv: R.Tensor((), dtype="float32") = R.einsum((inp_0,), subscripts="ii")
                gv: R.Tensor((), dtype="float32") = lv
                R.output(gv)
            return gv

    @tvm.script.ir_module
    class Expected2:
        @R.function
        def main(
            inp_0: R.Tensor((5,), dtype="float32"), inp_1: R.Tensor((4,), dtype="float32")
        ) -> R.Tensor((5, 4), dtype="float32"):
            with R.dataflow():
                lv: R.Tensor((5, 4), dtype="float32") = R.einsum(
                    (inp_0, inp_1), subscripts="i,j->ij"
                )
                gv: R.Tensor((5, 4), dtype="float32") = lv
                R.output(gv)
            return gv

    verify_model(Einsum1(), [([4, 4], "float32")], {}, Expected1)
    verify_model(Einsum2(), [([5], "float32"), ([4], "float32")], {}, Expected2)


def test_outer():
    class Outer(torch.nn.Module):
        def forward(self, x, y):
            return torch.outer(x, y)

    @tvm.script.ir_module
    class expected:
        @R.function
        def main(
            a: R.Tensor((3,), dtype="float32"), b: R.Tensor((4,), dtype="float32")
        ) -> R.Tensor((3, 4), dtype="float32"):
            with R.dataflow():
                lv: R.Tensor((3, 4), dtype="float32") = R.outer(a, b)
                gv: R.Tensor((3, 4), dtype="float32") = lv
                R.output(gv)
            return gv

    input_infos = [([3], "float32"), ([4], "float32")]
    verify_model(Outer(), input_infos, {}, expected)


@pytest.mark.gpu
@pytest.mark.skipif(not env.has_gpu(), reason="need gpu")
def test_softplus():
    import torch
    from torch.nn import Module

    torch.set_grad_enabled(False)

    class Softplus0(torch.nn.Module):
        def __init__(self):
            super().__init__()
            self.softplus = torch.nn.Softplus(1.0, 20.0)

        def forward(self, x):
            return self.softplus(x)

    class Softplus1(Module):
        def forward(self, input):
            return torch.nn.functional.softplus(input, 1.0, 20.0)

    @tvm.script.ir_module
    class expected:
        @R.function
        def main(inp_0: R.Tensor((10, 10), dtype="float32")) -> R.Tensor((10, 10), dtype="float32"):
            # block 0
            with R.dataflow():
                lv: R.Tensor((10, 10), dtype="float32") = R.nn.softplus(
                    inp_0, beta=1.0, threshold=20.0
                )
                gv: R.Tensor((10, 10), dtype="float32") = lv
                R.output(gv)
            return gv

    input_info = [([10, 10], "float32")]
    verify_model(Softplus0(), input_info, {}, expected)
    verify_model(Softplus1(), input_info, {}, expected)


@pytest.mark.gpu
@pytest.mark.skipif(not env.has_gpu(), reason="need gpu")
def test_leakyrelu():
    import torch
    from torch.nn import Module

    torch.set_grad_enabled(False)

    class LeakyReLU0(Module):
        def __init__(self):
            super().__init__()
            self.leakyrelu = torch.nn.LeakyReLU(0.02)

        def forward(self, input):
            return self.leakyrelu(input)

    class LeakyReLU1(Module):
        def forward(self, input):
            return torch.nn.functional.leaky_relu(input, 0.02)

    @tvm.script.ir_module
    class expected:
        @R.function
        def main(input_1: R.Tensor((10, 10), dtype="float32")) -> R.Tensor(
            (10, 10), dtype="float32"
        ):
            # block 0
            with R.dataflow():
                lv: R.Tensor((10, 10), dtype="float32") = R.nn.leakyrelu(input_1, 0.02)
                gv: R.Tensor((10, 10), dtype="float32") = lv
                R.output(gv)
            return gv

    input_info = [([10, 10], "float32")]
    verify_model(LeakyReLU0(), input_info, {}, expected)
    verify_model(LeakyReLU1(), input_info, {}, expected)


def test_prelu():
    class Prelu1(Module):
        def __init__(self, num_parameters=1, alpha=0.25):
            super().__init__()
            self.prelu = torch.nn.PReLU(num_parameters=num_parameters, init=alpha)

        def forward(self, x):
            return self.prelu(x)

    class Prelu2(torch.nn.Module):
        def __init__(self):
            super().__init__()
            self.alpha = torch.nn.Parameter(torch.tensor([0.25]))

        def forward(self, x):
            return torch.nn.functional.prelu(x, self.alpha)

    @tvm.script.ir_module
    class expected:
        @R.function
        def main(x: R.Tensor((1, 3, 10, 10), dtype="float32")) -> R.Tensor(
            (1, 3, 10, 10), dtype="float32"
        ):
            with R.dataflow():
                lv: R.Tensor((1, 3, 10, 10), dtype="float32") = R.nn.prelu(
                    x, R.const([0.25], dtype="float32"), axis=1
                )
                gv: R.Tensor((1, 3, 10, 10), dtype="float32") = lv
                R.output(gv)
            return gv

    input_info = [([1, 3, 10, 10], "float32")]
    verify_model(Prelu1(), input_info, {}, expected)
    verify_model(Prelu2(), input_info, {}, expected)


def test_maxpool1d():
    input_info = [([1, 3, 10], "float32")]

    class MaxPool1d(Module):
        def __init__(self):
            super().__init__()
            self.pool = torch.nn.MaxPool1d(kernel_size=2)

        def forward(self, input):
            return self.pool(input)

    class MaxPool1d_functional(Module):
        def __init__(self):
            super().__init__()

        def forward(self, input):
            return torch.nn.functional.max_pool1d(input, kernel_size=2)

    @tvm.script.ir_module
    class expected1:
        @R.function
        def main(input_1: R.Tensor((1, 3, 10), dtype="float32")) -> R.Tensor(
            (1, 3, 5), dtype="float32"
        ):
            with R.dataflow():
                lv: R.Tensor((1, 3, 5), dtype="float32") = R.nn.max_pool1d(
                    input_1,
                    pool_size=[2],
                    strides=[2],
                    dilation=[1],
                    padding=[0, 0],
                    layout="NCW",
                    out_layout="NCW",
                )
                gv: R.Tensor((1, 3, 5), dtype="float32") = lv
                R.output(gv)
            return gv

    class MaxPool1d2(Module):
        def __init__(self):
            super().__init__()
            self.pool = torch.nn.MaxPool1d(kernel_size=3, stride=1, padding=1)

        def forward(self, input):
            return self.pool(input)

    @tvm.script.ir_module
    class expected2:
        @R.function
        def main(input_1: R.Tensor((1, 3, 10), dtype="float32")) -> R.Tensor(
            (1, 3, 10), dtype="float32"
        ):
            with R.dataflow():
                lv: R.Tensor((1, 3, 10), dtype="float32") = R.nn.max_pool1d(
                    input_1,
                    pool_size=[3],
                    strides=[1],
                    dilation=[1],
                    padding=[1, 1],
                    layout="NCW",
                    out_layout="NCW",
                )
                gv: R.Tensor((1, 3, 10), dtype="float32") = lv
                R.output(gv)
            return gv

    class MaxPool1d3(Module):
        def __init__(self):
            super().__init__()
            self.pool = torch.nn.MaxPool1d(kernel_size=3, stride=2, dilation=2)

        def forward(self, input):
            return self.pool(input)

    @tvm.script.ir_module
    class expected3:
        @R.function
        def main(input_1: R.Tensor((1, 3, 10), dtype="float32")) -> R.Tensor(
            (1, 3, 3), dtype="float32"
        ):  # Corrected here
            with R.dataflow():
                lv: R.Tensor((1, 3, 3), dtype="float32") = R.nn.max_pool1d(
                    input_1,
                    pool_size=[3],
                    strides=[2],
                    dilation=[2],
                    padding=[0, 0],
                    layout="NCW",
                    out_layout="NCW",
                )
                gv: R.Tensor((1, 3, 3), dtype="float32") = lv
                R.output(gv)
            return gv

    verify_model(MaxPool1d(), input_info, {}, expected1)
    verify_model(MaxPool1d_functional(), input_info, {}, expected1)
    verify_model(MaxPool1d2(), input_info, {}, expected2)
    verify_model(MaxPool1d3(), input_info, {}, expected3)


def test_maxpool2d():
    input_info = [([1, 3, 10, 10], "float32")]

    class MaxPool2d(Module):
        def __init__(self):
            super().__init__()
            self.pool = torch.nn.MaxPool2d(kernel_size=[1, 1])

        def forward(self, input):
            return self.pool(input)

    class MaxPool2d_functional(Module):
        def __init__(self):
            super().__init__()

        def forward(self, input):
            return torch.nn.functional.max_pool2d(input, kernel_size=[1, 1])

    @tvm.script.ir_module
    class expected1:
        @R.function
        def main(input_1: R.Tensor((1, 3, 10, 10), dtype="float32")) -> R.Tensor(
            (1, 3, 10, 10), dtype="float32"
        ):
            # block 0
            with R.dataflow():
                lv: R.Tensor((1, 3, 10, 10), dtype="float32") = R.nn.max_pool2d(
                    input_1,
                    pool_size=[1, 1],
                    strides=[1, 1],
                    dilation=[1, 1],
                    padding=[0, 0, 0, 0],
                    layout="NCHW",
                    out_layout="NCHW",
                )
                gv: R.Tensor((1, 3, 10, 10), dtype="float32") = lv
                R.output(gv)
            return gv

    class MaxPool2d2(Module):
        def __init__(self):
            super().__init__()
            self.pool = torch.nn.MaxPool2d(kernel_size=[2, 2], dilation=[2, 3])

        def forward(self, input):
            return self.pool(input)

    @tvm.script.ir_module
    class expected2:
        @R.function
        def main(input_1: R.Tensor((1, 3, 10, 10), dtype="float32")) -> R.Tensor(
            (1, 3, 4, 4), dtype="float32"
        ):
            # block 0
            with R.dataflow():
                lv: R.Tensor((1, 3, 4, 4), dtype="float32") = R.nn.max_pool2d(
                    input_1,
                    pool_size=[2, 2],
                    strides=[2, 2],
                    dilation=[2, 3],
                    padding=[0, 0, 0, 0],
                    layout="NCHW",
                    out_layout="NCHW",
                )
                gv: R.Tensor((1, 3, 4, 4), dtype="float32") = lv
                R.output(gv)
            return gv

    class MaxPool2d3(Module):
        def __init__(self):
            super().__init__()
            self.pool = torch.nn.MaxPool2d(kernel_size=[4, 4], padding=2, stride=2)

        def forward(self, input):
            return self.pool(input)

    @tvm.script.ir_module
    class expected3:
        @R.function
        def main(input_1: R.Tensor((1, 3, 10, 10), dtype="float32")) -> R.Tensor(
            (1, 3, 6, 6), dtype="float32"
        ):
            # block 0
            with R.dataflow():
                lv: R.Tensor((1, 3, 6, 6), dtype="float32") = R.nn.max_pool2d(
                    input_1,
                    pool_size=[4, 4],
                    strides=[2, 2],
                    dilation=[1, 1],
                    padding=[2, 2, 2, 2],
                    layout="NCHW",
                    out_layout="NCHW",
                )
                gv: R.Tensor((1, 3, 6, 6), dtype="float32") = lv
                R.output(gv)
            return gv

    verify_model(MaxPool2d(), input_info, {}, expected1)
    verify_model(MaxPool2d_functional(), input_info, {}, expected1)
    verify_model(MaxPool2d2(), input_info, {}, expected2)
    verify_model(MaxPool2d3(), input_info, {}, expected3)


def test_maxpool3d():
    input_info = [([1, 3, 10, 10, 10], "float32")]

    class MaxPool3d(Module):
        def __init__(self):
            super().__init__()
            self.pool = torch.nn.MaxPool3d(kernel_size=[1, 1, 1])

        def forward(self, input):
            return self.pool(input)

    class MaxPool3d_functional(Module):
        def __init__(self):
            super().__init__()

        def forward(self, input):
            return torch.nn.functional.max_pool3d(input, kernel_size=[1, 1, 1])

    @tvm.script.ir_module
    class expected1:
        @R.function
        def main(input_1: R.Tensor((1, 3, 10, 10, 10), dtype="float32")) -> R.Tensor(
            (1, 3, 10, 10, 10), dtype="float32"
        ):
            with R.dataflow():
                lv: R.Tensor((1, 3, 10, 10, 10), dtype="float32") = R.nn.max_pool3d(
                    input_1,
                    pool_size=[1, 1, 1],
                    strides=[1, 1, 1],
                    dilation=[1, 1, 1],
                    padding=[0, 0, 0, 0, 0, 0],
                    layout="NCDHW",
                    out_layout="NCDHW",
                )
                gv: R.Tensor((1, 3, 10, 10, 10), dtype="float32") = lv
                R.output(gv)
            return gv

    class MaxPool3d2(Module):
        def __init__(self):
            super().__init__()
            self.pool = torch.nn.MaxPool3d(kernel_size=[2, 2, 2], dilation=[1, 2, 2])

        def forward(self, input):
            return self.pool(input)

    @tvm.script.ir_module
    class expected2:
        @R.function
        def main(input_1: R.Tensor((1, 3, 10, 10, 10), dtype="float32")) -> R.Tensor(
            (1, 3, 5, 4, 4), dtype="float32"
        ):  # Fixed here
            with R.dataflow():
                lv: R.Tensor((1, 3, 5, 4, 4), dtype="float32") = R.nn.max_pool3d(
                    input_1,
                    pool_size=[2, 2, 2],
                    strides=[2, 2, 2],
                    dilation=[1, 2, 2],
                    padding=[0, 0, 0, 0, 0, 0],
                    layout="NCDHW",
                    out_layout="NCDHW",
                )
                gv: R.Tensor((1, 3, 5, 4, 4), dtype="float32") = lv
                R.output(gv)
            return gv

    class MaxPool3d3(Module):
        def __init__(self):
            super().__init__()
            self.pool = torch.nn.MaxPool3d(kernel_size=[3, 3, 3], padding=1, stride=2)

        def forward(self, input):
            return self.pool(input)

    @tvm.script.ir_module
    class expected3:
        @R.function
        def main(input_1: R.Tensor((1, 3, 10, 10, 10), dtype="float32")) -> R.Tensor(
            (1, 3, 5, 5, 5), dtype="float32"
        ):
            with R.dataflow():
                lv: R.Tensor((1, 3, 5, 5, 5), dtype="float32") = R.nn.max_pool3d(
                    input_1,
                    pool_size=[3, 3, 3],
                    strides=[2, 2, 2],
                    dilation=[1, 1, 1],
                    padding=[1, 1, 1, 1, 1, 1],
                    layout="NCDHW",
                    out_layout="NCDHW",
                )
                gv: R.Tensor((1, 3, 5, 5, 5), dtype="float32") = lv
                R.output(gv)
            return gv

    verify_model(MaxPool3d(), input_info, {}, expected1)
    verify_model(MaxPool3d_functional(), input_info, {}, expected1)
    verify_model(MaxPool3d2(), input_info, {}, expected2)
    verify_model(MaxPool3d3(), input_info, {}, expected3)


def test_avgpool1d():
    input_info = [([1, 3, 10], "float32")]

    class AvgPool1d(Module):
        def __init__(self):
            super().__init__()
            self.pool = torch.nn.AvgPool1d(kernel_size=1)

        def forward(self, input):
            return self.pool(input)

    @tvm.script.ir_module
    class expected1:
        @R.function
        def main(input_1: R.Tensor((1, 3, 10), dtype="float32")) -> R.Tensor(
            (1, 3, 10), dtype="float32"
        ):
            with R.dataflow():
                lv = R.nn.avg_pool1d(
                    input_1,
                    pool_size=[1],
                    strides=[1],
                    dilation=[1],
                    padding=[0, 0],
                    ceil_mode=False,
                    count_include_pad=True,
                    layout="NCW",
                    out_layout="NCW",
                )
                gv = lv
                R.output(gv)
            return gv

    class AvgPool1d2(Module):
        def __init__(self):
            super().__init__()
            self.pool = torch.nn.AvgPool1d(kernel_size=4, stride=2, padding=2, ceil_mode=True)

        def forward(self, input):
            return self.pool(input)

    class AvgPool1d3(Module):
        def forward(self, input):
            return torch.nn.functional.avg_pool1d(
                input, kernel_size=4, stride=2, padding=2, ceil_mode=True
            )

    @tvm.script.ir_module
    class expected2:
        @R.function
        def main(input_1: R.Tensor((1, 3, 10), dtype="float32")):
            with R.dataflow():
                lv = R.nn.avg_pool1d(
                    input_1,
                    pool_size=[4],
                    strides=[2],
                    dilation=[1],
                    padding=[2, 2],
                    ceil_mode=True,
                    count_include_pad=True,
                    layout="NCW",
                    out_layout="NCW",
                )
                gv = lv
                R.output(gv)
            return gv

    class AvgPool1d4(Module):
        def forward(self, input):
            return torch.nn.functional.avg_pool1d(input, kernel_size=2)

    @tvm.script.ir_module
    class expected3:
        @R.function
        def main(input_1: R.Tensor((1, 3, 10), dtype="float32")):
            with R.dataflow():
                lv = R.nn.avg_pool1d(
                    input_1,
                    pool_size=[2],
                    strides=[2],
                    dilation=[1],
                    padding=[0, 0],
                    ceil_mode=False,
                    count_include_pad=True,
                    layout="NCW",
                    out_layout="NCW",
                )
                gv = lv
                R.output(gv)
            return gv

    verify_model(AvgPool1d(), input_info, {}, expected1)
    verify_model(AvgPool1d2(), input_info, {}, expected2)
    verify_model(AvgPool1d3(), input_info, {}, expected2)
    verify_model(AvgPool1d4(), input_info, {}, expected3)


def test_avgpool2d():
    input_info = [([1, 3, 10, 10], "float32")]

    class AvgPool2d(Module):
        def __init__(self):
            super().__init__()
            self.pool = torch.nn.AvgPool2d(kernel_size=[1, 1])

        def forward(self, input):
            return self.pool(input)

    @tvm.script.ir_module
    class expected1:
        @R.function
        def main(input_1: R.Tensor((1, 3, 10, 10), dtype="float32")) -> R.Tensor(
            (1, 3, 10, 10), dtype="float32"
        ):
            # block 0
            with R.dataflow():
                lv: R.Tensor((1, 3, 10, 10), dtype="float32") = R.nn.avg_pool2d(
                    input_1,
                    pool_size=[1, 1],
                    strides=[1, 1],
                    dilation=[1, 1],
                    padding=[0, 0, 0, 0],
                    count_include_pad=True,
                    layout="NCHW",
                    out_layout="NCHW",
                )
                gv: R.Tensor((1, 3, 10, 10), dtype="float32") = lv
                R.output(gv)
            return gv

    class AvgPool2d2(Module):
        def __init__(self):
            super().__init__()
            self.pool = torch.nn.AvgPool2d(kernel_size=[4, 4], stride=2, padding=2, ceil_mode=True)

        def forward(self, input):
            return self.pool(input)

    class AvgPool2d3(Module):
        def forward(self, input):
            return torch.nn.functional.avg_pool2d(
                input, kernel_size=[4, 4], stride=2, padding=2, ceil_mode=True
            )

    @tvm.script.ir_module
    class expected2:
        @R.function
        def main(input_1: R.Tensor((1, 3, 10, 10), dtype="float32")):
            with R.dataflow():
                lv = R.nn.avg_pool2d(
                    input_1,
                    pool_size=[4, 4],
                    strides=[2, 2],
                    dilation=[1, 1],
                    padding=[2, 2, 2, 2],
                    ceil_mode=True,
                    count_include_pad=True,
                    layout="NCHW",
                    out_layout="NCHW",
                )
                gv = lv
                R.output(gv)
            return gv

    class AvgPool2d4(Module):
        def forward(self, input):
            return torch.nn.functional.avg_pool2d(input, kernel_size=[2, 1], divisor_override=2)

    @tvm.script.ir_module
    class expected3:
        @R.function
        def main(input_1: R.Tensor((1, 3, 10, 10), dtype="float32")):
            with R.dataflow():
                lv = R.nn.avg_pool2d(
                    input_1,
                    pool_size=[2, 1],
                    strides=[2, 1],
                    dilation=[1, 1],
                    padding=[0, 0, 0, 0],
                    ceil_mode=False,
                    count_include_pad=True,
                    layout="NCHW",
                    out_layout="NCHW",
                )
                gv = lv
                R.output(gv)
            return gv

    verify_model(AvgPool2d(), input_info, {}, expected1)
    verify_model(AvgPool2d2(), input_info, {}, expected2)
    verify_model(AvgPool2d3(), input_info, {}, expected2)
    verify_model(AvgPool2d4(), input_info, {}, expected3)


def test_avgpool3d():
    input_info = [([1, 3, 8, 8, 8], "float32")]

    class AvgPool3d(Module):
        def __init__(self):
            super().__init__()
            self.pool = torch.nn.AvgPool3d(kernel_size=[1, 1, 1])

        def forward(self, input):
            return self.pool(input)

    @tvm.script.ir_module
    class expected1:
        @R.function
        def main(input_1: R.Tensor((1, 3, 8, 8, 8), dtype="float32")) -> R.Tensor(
            (1, 3, 8, 8, 8), dtype="float32"
        ):
            with R.dataflow():
                lv: R.Tensor((1, 3, 8, 8, 8), dtype="float32") = R.nn.avg_pool3d(
                    input_1,
                    pool_size=[1, 1, 1],
                    strides=[1, 1, 1],
                    dilation=[1, 1, 1],
                    padding=[0, 0, 0, 0, 0, 0],
                    ceil_mode=False,
                    count_include_pad=True,
                    layout="NCDHW",
                    out_layout="NCDHW",
                )
                gv: R.Tensor((1, 3, 8, 8, 8), dtype="float32") = lv
                R.output(gv)
            return gv

    class AvgPool3d2(Module):
        def __init__(self):
            super().__init__()
            self.pool = torch.nn.AvgPool3d(
                kernel_size=[3, 3, 3], stride=2, padding=1, ceil_mode=True
            )

        def forward(self, input):
            return self.pool(input)

    class AvgPool3d3(Module):
        def forward(self, input):
            return torch.nn.functional.avg_pool3d(
                input, kernel_size=[3, 3, 3], stride=2, padding=1, ceil_mode=True
            )

    @tvm.script.ir_module
    class expected2:
        @R.function
        def main(input_1: R.Tensor((1, 3, 8, 8, 8), dtype="float32")):
            with R.dataflow():
                lv = R.nn.avg_pool3d(
                    input_1,
                    pool_size=[3, 3, 3],
                    strides=[2, 2, 2],
                    dilation=[1, 1, 1],
                    padding=[1, 1, 1, 1, 1, 1],
                    ceil_mode=True,
                    count_include_pad=True,
                    layout="NCDHW",
                    out_layout="NCDHW",
                )
                gv = lv
                R.output(gv)
            return gv

    class AvgPool3d4(Module):
        def forward(self, input):
            return torch.nn.functional.avg_pool3d(input, kernel_size=[2, 1, 2], stride=[2, 1, 2])

    @tvm.script.ir_module
    class expected3:
        @R.function
        def main(input_1: R.Tensor((1, 3, 8, 8, 8), dtype="float32")):
            with R.dataflow():
                lv = R.nn.avg_pool3d(
                    input_1,
                    pool_size=[2, 1, 2],
                    strides=[2, 1, 2],
                    dilation=[1, 1, 1],
                    padding=[0, 0, 0, 0, 0, 0],
                    ceil_mode=False,
                    count_include_pad=True,
                    layout="NCDHW",
                    out_layout="NCDHW",
                )
                gv = lv
                R.output(gv)
            return gv

    verify_model(AvgPool3d(), input_info, {}, expected1)
    verify_model(AvgPool3d2(), input_info, {}, expected2)
    verify_model(AvgPool3d3(), input_info, {}, expected2)
    verify_model(AvgPool3d4(), input_info, {}, expected3)


def test_adaptive_avgpool1d():
    input_info = [([1, 3, 16], "float32")]

    class AdaptiveAvgPool1d0(torch.nn.Module):
        def __init__(self):
            super().__init__()
            self.pool = torch.nn.AdaptiveAvgPool1d(8)

        def forward(self, input):
            return self.pool(input)

    class AdaptiveAvgPool1d1(torch.nn.Module):
        def forward(self, input):
            return torch.nn.functional.adaptive_avg_pool1d(input, 8)

    @tvm.script.ir_module
    class expected1:
        @R.function
        def main(input_1: R.Tensor((1, 3, 16), dtype="float32")) -> R.Tensor(
            (1, 3, 8), dtype="float32"
        ):
            with R.dataflow():
                lv: R.Tensor((1, 3, 8), dtype="float32") = R.nn.adaptive_avg_pool1d(
                    input_1, output_size=[8], layout="NCW", out_layout="NCW"
                )
                gv: R.Tensor((1, 3, 8), dtype="float32") = lv
                R.output(gv)
            return gv

    verify_model(AdaptiveAvgPool1d0(), input_info, {}, expected1)
    verify_model(AdaptiveAvgPool1d1(), input_info, {}, expected1)


def test_adaptive_avgpool2d():
    input_info = [([1, 3, 10, 10], "float32")]

    class AdaptiveAvgPool2d0(Module):
        def __init__(self):
            super().__init__()
            self.pool = torch.nn.AdaptiveAvgPool2d([10, 10])

        def forward(self, input):
            return self.pool(input)

    class AdaptiveAvgPool2d1(Module):
        def forward(self, input):
            return torch.nn.functional.adaptive_avg_pool2d(input, [10, 10])

    @tvm.script.ir_module
    class expected1:
        @R.function
        def main(input_1: R.Tensor((1, 3, 10, 10), dtype="float32")) -> R.Tensor(
            (1, 3, 10, 10), dtype="float32"
        ):
            # block 0
            with R.dataflow():
                lv: R.Tensor((1, 3, 10, 10), dtype="float32") = R.nn.adaptive_avg_pool2d(
                    input_1, output_size=[10, 10], layout="NCHW", out_layout="NCHW"
                )
                gv: R.Tensor((1, 3, 10, 10), dtype="float32") = lv
                R.output(gv)
            return gv

    verify_model(AdaptiveAvgPool2d0(), input_info, {}, expected1)
    verify_model(AdaptiveAvgPool2d1(), input_info, {}, expected1)


def test_adaptive_avgpool3d():
    input_info = [([1, 3, 16, 16, 16], "float32")]

    class AdaptiveAvgPool3d0(torch.nn.Module):
        def __init__(self):
            super().__init__()
            self.pool = torch.nn.AdaptiveAvgPool3d((8, 8, 8))

        def forward(self, input):
            return self.pool(input)

    class AdaptiveAvgPool3d1(torch.nn.Module):
        def forward(self, input):
            return torch.nn.functional.adaptive_avg_pool3d(input, (8, 8, 8))

    @tvm.script.ir_module
    class expected1:
        @R.function
        def main(input_1: R.Tensor((1, 3, 16, 16, 16), dtype="float32")) -> R.Tensor(
            (1, 3, 8, 8, 8), dtype="float32"
        ):
            with R.dataflow():
                lv: R.Tensor((1, 3, 8, 8, 8), dtype="float32") = R.nn.adaptive_avg_pool3d(
                    input_1, output_size=[8, 8, 8], layout="NCDHW", out_layout="NCDHW"
                )
                gv: R.Tensor((1, 3, 8, 8, 8), dtype="float32") = lv
                R.output(gv)
            return gv

    verify_model(AdaptiveAvgPool3d0(), input_info, {}, expected1)
    verify_model(AdaptiveAvgPool3d1(), input_info, {}, expected1)


def test_flatten():
    input_info = [([1, 3, 10, 10], "float32")]

    class Flatten(Module):
        def __init__(self):
            super().__init__()
            self.f = torch.nn.Flatten(2, -1)

        def forward(self, input):
            return self.f(input)

    @tvm.script.ir_module
    class expected1:
        @R.function
        def main(input_1: R.Tensor((1, 3, 10, 10), dtype="float32")) -> R.Tensor(
            (1, 3, 100), dtype="float32"
        ):
            # block 0
            with R.dataflow():
                lv: R.Tensor((1, 3, 100), dtype="float32") = R.reshape(input_1, (1, 3, 100))
                gv: R.Tensor((1, 3, 100), dtype="float32") = lv
                R.output(gv)
            return gv

    # call_module
    verify_model(Flatten(), input_info, {}, expected1)
    # call_method
    verify_model(torch.nn.Flatten(2, -1), input_info, {}, expected1)


def test_batchnorm2d():
    input_info = [([1, 3, 10, 10], "float32")]

    class BatchNorm2d(Module):
        def __init__(self):
            super().__init__()
            self.bn = torch.nn.BatchNorm2d(3)

        def forward(self, input):
            return self.bn(input)

    @tvm.script.ir_module
    class expected1:
        @R.function
        def main(
            input_1: R.Tensor((1, 3, 10, 10), dtype="float32"),
            w1: R.Tensor((3,), dtype="float32"),
            w2: R.Tensor((3,), dtype="float32"),
            w3: R.Tensor((3,), dtype="float32"),
            w4: R.Tensor((3,), dtype="float32"),
        ) -> R.Tensor((1, 3, 10, 10), dtype="float32"):
            # block 0
            with R.dataflow():
                lv: R.Tuple(
                    R.Tensor((1, 3, 10, 10), dtype="float32"),
                    R.Tensor((3,), dtype="float32"),
                    R.Tensor((3,), dtype="float32"),
                ) = R.nn.batch_norm(
                    input_1,
                    w1,
                    w2,
                    w3,
                    w4,
                    axis=1,
                    epsilon=1e-05,
                    center=True,
                    scale=True,
                )
                lv1: R.Tensor((1, 3, 10, 10), dtype="float32") = lv[0]
                gv: R.Tensor((1, 3, 10, 10), dtype="float32") = lv1
                R.output(gv)
            return gv

    model = BatchNorm2d()
    binding = {
        "w1": model.bn.weight.detach().numpy(),
        "w2": model.bn.bias.detach().numpy(),
        "w3": model.bn.running_mean.detach().numpy(),
        "w4": model.bn.running_var.detach().numpy(),
    }
    verify_model(BatchNorm2d(), input_info, binding, expected1)


def test_embedding():
    input_info = [([4], "int64")]

    class Embedding(Module):
        def __init__(self):
            super().__init__()
            self.embedding = torch.nn.Embedding(10, 3)

        def forward(self, input):
            return self.embedding(input)

    @tvm.script.ir_module
    class expected1:
        @R.function
        def main(
            input_1: R.Tensor((4,), dtype="int64"), w1: R.Tensor((10, 3), dtype="float32")
        ) -> R.Tensor((4, 3), dtype="float32"):
            # block 0
            with R.dataflow():
                lv: R.Tensor((4,), dtype="int32") = R.astype(input_1, dtype="int32")
                lv1: R.Tensor((4, 3), dtype="float32") = R.take(w1, lv, axis=0)
                gv: R.Tensor((4, 3), dtype="float32") = lv1
                R.output(gv)
            return gv

    model = Embedding()
    binding = {"w1": model.embedding.weight.detach().numpy()}
    verify_model(model, input_info, binding, expected1)


def test_stochastic_depth():
    torchvision = pytest.importorskip("torchvision")

    input_info = [([1, 3, 10, 10], "float32")]

    class StochasticDepth1(Module):
        def __init__(self):
            super().__init__()
            self.stochastic_depth = torchvision.ops.StochasticDepth(0.5, mode="row")

        def forward(self, x):
            return self.stochastic_depth(x)

    class StochasticDepth2(Module):
        def forward(self, x):
            return torchvision.ops.stochastic_depth(x, 0.5, mode="row", training=False)

    @tvm.script.ir_module
    class expected1:
        @R.function
        def main(input_1: R.Tensor((1, 3, 10, 10), dtype="float32")) -> R.Tensor(
            (1, 3, 10, 10), dtype="float32"
        ):
            # block 0
            with R.dataflow():
                gv: R.Tensor((1, 3, 10, 10), dtype="float32") = input_1
                R.output(gv)
            return gv

    verify_model(StochasticDepth1(), input_info, {}, expected1)
    verify_model(StochasticDepth2(), input_info, {}, expected1)


def test_layernorm():
    input_info = [([1, 3, 10, 10], "float32")]

    class LayerNorm(Module):
        def __init__(self):
            super().__init__()
            self.ln = torch.nn.LayerNorm((10, 10))

        def forward(self, input):
            return self.ln(input)

    @tvm.script.ir_module
    class expected1:
        @R.function
        def main(
            input_1: R.Tensor((1, 3, 10, 10), dtype="float32"),
            w1: R.Tensor((10, 10), dtype="float32"),
            w2: R.Tensor((10, 10), dtype="float32"),
        ) -> R.Tensor((1, 3, 10, 10), dtype="float32"):
            # block 0
            with R.dataflow():
                lv: R.Tensor((1, 3, 10, 10), dtype="float32") = R.nn.layer_norm(
                    input_1,
                    w1,
                    w2,
                    axes=[-2, -1],
                    epsilon=1e-05,
                    center=True,
                    scale=True,
                )
                gv: R.Tensor((1, 3, 10, 10), dtype="float32") = lv
                R.output(gv)
            return gv

    model = LayerNorm()
    binding = {
        "w1": model.ln.weight.detach().numpy(),
        "w2": model.ln.bias.detach().numpy(),
    }
    verify_model(LayerNorm(), input_info, binding, expected1)


def test_functional_layernorm():
    import numpy as np

    input_info = [([1, 3, 10, 10], "float32")]

    class LayerNorm(Module):
        def __init__(self, shape):
            super().__init__()
            self.weight = torch.nn.Parameter(torch.ones(shape))
            self.bias = torch.nn.Parameter(torch.zeros(shape))

        def forward(self, input):
            return torch.nn.functional.layer_norm(
                input, self.weight.shape, self.weight, self.bias, 1e-5
            )

    @tvm.script.ir_module
    class expected1:
        @R.function
        def main(
            input_1: R.Tensor((1, 3, 10, 10), dtype="float32"),
            w1: R.Tensor((10, 10), dtype="float32"),
            w2: R.Tensor((10, 10), dtype="float32"),
        ) -> R.Tensor((1, 3, 10, 10), dtype="float32"):
            # block 0
            with R.dataflow():
                lv: R.Tensor((1, 3, 10, 10), dtype="float32") = R.nn.layer_norm(
                    input_1,
                    w1,
                    w2,
                    axes=[-2, -1],
                    epsilon=1e-05,
                    center=True,
                    scale=True,
                )
                gv: R.Tensor((1, 3, 10, 10), dtype="float32") = lv
                R.output(gv)
            return gv

    model = LayerNorm((10, 10))
    binding = {
        "w1": model.weight.detach().numpy(),
        "w2": model.bias.detach().numpy(),
    }
    verify_model(model, input_info, binding, expected1)

    class LayerNorm2(Module):
        def __init__(self, shape):
            super().__init__()
            self.shape = shape
            self.weight = None
            self.bias = None

        def forward(self, input):
            return torch.nn.functional.layer_norm(input, self.shape, self.weight, self.bias, 1e-5)

    @tvm.script.ir_module
    class expected2:
        @R.function
        def main(
            input_1: R.Tensor((1, 3, 10, 10), dtype="float32"),
        ) -> R.Tensor((1, 3, 10, 10), dtype="float32"):
            with R.dataflow():
                lv: R.Tensor((1, 3, 10, 10), dtype="float32") = R.nn.layer_norm(
                    input_1,
                    gamma=relax.const(np.ones((10, 10)), dtype="float32"),
                    beta=relax.const(np.zeros((10, 10)), dtype="float32"),
                    axes=[-2, -1],
                    epsilon=1e-05,
                    center=True,
                    scale=True,
                )
                gv: R.Tensor((1, 3, 10, 10), dtype="float32") = lv
                R.output(gv)
            return gv

    model = LayerNorm2((10, 10))
    binding = {}
    verify_model(model, input_info, binding, expected2)

    class LayerNorm3(Module):
        def __init__(self, shape):
            super().__init__()
            self.shape = shape
            self.weight = torch.nn.Parameter(torch.ones(shape))
            self.bias = torch.nn.Parameter(torch.zeros(shape))

        def forward(self, input):
            return torch.nn.functional.layer_norm(input, self.shape, self.weight, self.bias, 1e-5)

    @tvm.script.ir_module
    class expected3:
        @R.function
        def main(
            input_1: R.Tensor((1, 3, 10, 10), dtype="float32"),
            w1: R.Tensor([10, 10], dtype="float32"),
            w2: R.Tensor([10, 10], dtype="float32"),
        ) -> R.Tensor((1, 3, 10, 10), dtype="float32"):
            # block 0
            with R.dataflow():
                lv: R.Tensor((1, 3, 10, 10), dtype="float32") = R.nn.layer_norm(
                    input_1,
                    w1,
                    w2,
                    axes=[-2, -1],
                    epsilon=1e-05,
                    center=True,
                    scale=True,
                )
                gv: R.Tensor((1, 3, 10, 10), dtype="float32") = lv
                R.output(gv)
            return gv

    model = LayerNorm3([10, 10])
    binding = {
        "w1": model.weight.detach().numpy(),
        "w2": model.bias.detach().numpy(),
    }
    verify_model(model, input_info, binding, expected3)


def test_cross_entropy():
    input_info = [([3, 2], "float32"), ([3], "int32")]

    class CrossEntropy1(Module):
        def __init__(self):
            super().__init__()
            self.loss = torch.nn.CrossEntropyLoss()

        def forward(self, logits, targets):
            return self.loss(logits, targets)

    @tvm.script.ir_module
    class expected1:
        @R.function
        def main(
            inp_0: R.Tensor((3, 2), dtype="float32"), inp_1: R.Tensor((3,), dtype="int32")
        ) -> R.Tensor((), dtype="float32"):
            with R.dataflow():
                lv: R.Tensor((3, 2), dtype="float32") = R.nn.log_softmax(inp_0, axis=-1)
                lv1: R.Tensor((), dtype="float32") = R.nn.nll_loss(
                    lv, inp_1, reduction="mean", ignore_index=-100
                )
                gv: R.Tensor((), dtype="float32") = lv1
                R.output(gv)
            return gv

    class CrossEntropy2(Module):
        def __init__(self):
            super().__init__()
            self.weight = torch.nn.Parameter(torch.ones((2,)))
            self.loss = torch.nn.CrossEntropyLoss(weight=self.weight)

        def forward(self, logits, targets):
            return self.loss(logits, targets)

    @tvm.script.ir_module
    class expected2:
        @R.function
        def main(
            inp_0: R.Tensor((3, 2), dtype="float32"),
            inp_1: R.Tensor((3,), dtype="int32"),
            w1: R.Tensor((2,), dtype="float32"),
        ) -> R.Tensor((), dtype="float32"):
            with R.dataflow():
                lv: R.Tensor((3, 2), dtype="float32") = R.nn.log_softmax(inp_0, axis=-1)
                lv1: R.Tensor((), dtype="float32") = R.nn.nll_loss(
                    lv,
                    inp_1,
                    w1,
                    reduction="mean",
                    ignore_index=-100,
                )
                gv: R.Tensor((), dtype="float32") = lv1
                R.output(gv)
            return gv

    class CrossEntropy3(Module):
        def __init__(self):
            super().__init__()
            self.loss = torch.nn.CrossEntropyLoss(ignore_index=1, reduction="sum")

        def forward(self, logits, targets):
            return self.loss(logits, targets)

    @tvm.script.ir_module
    class expected3:
        @R.function
        def main(
            inp_0: R.Tensor((3, 2), dtype="float32"), inp_1: R.Tensor((3,), dtype="int32")
        ) -> R.Tensor((), dtype="float32"):
            with R.dataflow():
                lv: R.Tensor((3, 2), dtype="float32") = R.nn.log_softmax(inp_0, axis=-1)
                lv1: R.Tensor((), dtype="float32") = R.nn.nll_loss(
                    lv, inp_1, reduction="sum", ignore_index=1
                )
                gv: R.Tensor((), dtype="float32") = lv1
                R.output(gv)
            return gv

    verify_model(CrossEntropy1(), input_info, {}, expected1)
    model = CrossEntropy2()
    binding = {"w1": model.loss.weight.detach().numpy()}
    verify_model(model, input_info, binding, expected2)
    verify_model(CrossEntropy3(), input_info, {}, expected3)


def test_functional_cross_entropy():
    input_info = [([3, 10], "float32"), ([3], "int32")]

    class CrossEntropy(Module):
        def forward(self, logits, targets):
            return torch.nn.functional.cross_entropy(logits, targets)

    @tvm.script.ir_module
    class expected1:
        @R.function
        def main(
            inp_0: R.Tensor((3, 10), dtype="float32"), inp_1: R.Tensor((3,), dtype="int32")
        ) -> R.Tensor((), dtype="float32"):
            with R.dataflow():
                lv: R.Tensor((3, 10), dtype="float32") = R.nn.log_softmax(inp_0, axis=-1)
                lv1: R.Tensor((), dtype="float32") = R.nn.nll_loss(
                    lv, inp_1, reduction="mean", ignore_index=-100
                )
                gv: R.Tensor((), dtype="float32") = lv1
                R.output(gv)
            return gv

    model = CrossEntropy()
    verify_model(model, input_info, {}, expected1)


def test_groupnorm():
    import torch
    from torch.nn import Module

    torch.set_grad_enabled(False)
    torch.random.manual_seed(0)

    input_info = [([1, 3, 10, 10], "float32")]

    class GroupNorm(Module):
        def __init__(self):
            super().__init__()
            self.gn = torch.nn.GroupNorm(3, 3)

        def forward(self, input):
            return self.gn(input)

    @tvm.script.ir_module
    class expected1:
        @R.function
        def main(
            input_1: R.Tensor((1, 3, 10, 10), dtype="float32"),
            w1: R.Tensor((3,), dtype="float32"),
            w2: R.Tensor((3,), dtype="float32"),
        ) -> R.Tensor((1, 3, 10, 10), dtype="float32"):
            with R.dataflow():
                lv: R.Tensor((1, 3, 10, 10), dtype="float32") = R.nn.group_norm(
                    input_1,
                    w1,
                    w2,
                    num_groups=3,
                    channel_axis=1,
                    axes=[2, 3],
                    epsilon=1.0000000000000001e-05,
                    center=True,
                    scale=True,
                )
                gv: R.Tensor((1, 3, 10, 10), dtype="float32") = lv
                R.output(gv)
            return gv

    model = GroupNorm()
    binding = {
        "w1": model.gn.weight.detach().numpy(),
        "w2": model.gn.bias.detach().numpy(),
    }
    verify_model(model, input_info, binding, expected1)


def test_instancenorm2d():
    torch.set_grad_enabled(False)
    torch.random.manual_seed(0)

    input_info = [([1, 3, 10, 10], "float32")]

    class InstanceNorm2d(Module):
        def __init__(self):
            super().__init__()
            self.gn = torch.nn.InstanceNorm2d(3)

        def forward(self, input):
            return self.gn(input)

    @tvm.script.ir_module
    class expected1:
        @R.function
        def main(
            input_1: R.Tensor((1, 3, 10, 10), dtype="float32"),
            w1: R.Tensor((3,), dtype="float32"),
            w2: R.Tensor((3,), dtype="float32"),
        ) -> R.Tensor((1, 3, 10, 10), dtype="float32"):
            with R.dataflow():
                lv: R.Tensor((1, 3, 10, 10), dtype="float32") = R.nn.instance_norm(
                    input_1,
                    w1,
                    w2,
                    channel_axis=1,
                    axes=[2, 3],
                    epsilon=1e-05,
                    center=True,
                    scale=True,
                )
                gv: R.Tensor((1, 3, 10, 10), dtype="float32") = lv
                R.output(gv)
            return gv

    example_args = (torch.randn(1, 3, 10, 10, dtype=torch.float32),)

    model = InstanceNorm2d()
    binding = {
        "w1": torch.ones(3).detach().numpy(),
        "w2": torch.zeros(3).detach().numpy(),
    }
    verify_model(model, input_info, binding, expected1)


operator_binary_1 = [
    (operator.add, R.add),
    (operator.sub, R.subtract),
    (operator.mul, R.multiply),
    (operator.truediv, R.divide),
    (operator.floordiv, R.floor_divide),
    (torch.ops.aten.fmod, R.mod),
    (operator.pow, R.power),
    (operator.mod, R.floor_mod),
]


@pytest.mark.parametrize("op, relax_op", operator_binary_1)
def test_binary1(op, relax_op):
    input_info1 = [([1, 3, 10, 10], "float32"), ([1, 3, 10, 10], "float32")]
    input_info2 = [([1, 3, 10, 10], "float32")]

    class Binary1(Module):
        def __init__(self, op):
            super().__init__()
            self.op = op

        def forward(self, lhs, rhs):
            return self.op(lhs, rhs)

    @tvm.script.ir_module
    class expected_binary1:
        @R.function
        def main(
            lhs: R.Tensor((1, 3, 10, 10), dtype="float32"),
            rhs: R.Tensor((1, 3, 10, 10), dtype="float32"),
        ) -> R.Tensor((1, 3, 10, 10), dtype="float32"):
            with R.dataflow():
                lv: R.Tensor((1, 3, 10, 10), dtype="float32") = relax_op(lhs, rhs)
                gv: R.Tensor((1, 3, 10, 10), dtype="float32") = lv
                R.output(gv)
            return gv

    class Binary2(Module):
        def __init__(self, op):
            super().__init__()
            self.op = op

        def forward(self, lhs):
            return self.op(lhs, 1.0)

    @tvm.script.ir_module
    class expected_binary2:
        @R.function
        def main(
            lhs: R.Tensor((1, 3, 10, 10), dtype="float32"),
        ) -> R.Tensor((1, 3, 10, 10), dtype="float32"):
            with R.dataflow():
                lv: R.Tensor((1, 3, 10, 10), dtype="float32") = relax_op(lhs, R.const(1.0))
                gv: R.Tensor((1, 3, 10, 10), dtype="float32") = lv
                R.output(gv)
            return gv

    verify_model(Binary1(op), input_info1, {}, expected_binary1)
    verify_model(Binary2(op), input_info2, {}, expected_binary2)


operator_binary_2 = [
    (operator.eq, R.equal),
    (operator.ne, R.not_equal),
    (operator.lt, R.less),
    (operator.le, R.less_equal),
    (operator.gt, R.greater),
    (operator.ge, R.greater_equal),
]


@pytest.mark.parametrize("op, relax_op", operator_binary_2)
def test_binary2(op, relax_op):
    input_info1 = [([1, 3, 10, 10], "float32"), ([1, 3, 10, 10], "float32")]
    input_info2 = [([1, 3, 10, 10], "float32")]

    class Binary1(Module):
        def __init__(self, op):
            super().__init__()
            self.op = op

        def forward(self, lhs, rhs):
            return self.op(lhs, rhs)

    @tvm.script.ir_module
    class expected_binary1:
        @R.function
        def main(
            lhs: R.Tensor((1, 3, 10, 10), dtype="float32"),
            rhs: R.Tensor((1, 3, 10, 10), dtype="float32"),
        ) -> R.Tensor((1, 3, 10, 10), dtype="bool"):
            with R.dataflow():
                lv: R.Tensor((1, 3, 10, 10), dtype="bool") = relax_op(lhs, rhs)
                gv: R.Tensor((1, 3, 10, 10), dtype="bool") = lv
                R.output(gv)
            return gv

    class Binary2(Module):
        def __init__(self, op):
            super().__init__()
            self.op = op

        def forward(self, lhs):
            return self.op(lhs, 1.0)

    @tvm.script.ir_module
    class expected_binary2:
        @R.function
        def main(
            lhs: R.Tensor((1, 3, 10, 10), dtype="float32"),
        ) -> R.Tensor((1, 3, 10, 10), dtype="bool"):
            with R.dataflow():
                lv: R.Tensor((1, 3, 10, 10), dtype="bool") = relax_op(lhs, R.const(1.0))
                gv: R.Tensor((1, 3, 10, 10), dtype="bool") = lv
                R.output(gv)
            return gv

    verify_model(Binary1(op), input_info1, {}, expected_binary1)
    verify_model(Binary2(op), input_info2, {}, expected_binary2)


operator_binary_3 = [
    (torch.ops.aten.bitwise_or_, R.bitwise_or),
    (torch.ops.aten.bitwise_or, R.bitwise_or),
    (operator.lshift, R.left_shift),
    (operator.rshift, R.right_shift),
    (operator.and_, R.bitwise_and),
    (operator.or_, R.bitwise_or),
    (operator.xor, R.bitwise_xor),
]


@pytest.mark.parametrize("op, relax_op", operator_binary_3)
def test_binary3(op, relax_op):
    input_info1 = [([1, 3, 10, 10], "int32"), ([1, 3, 10, 10], "int32")]
    input_info2 = [([1, 3, 10, 10], "int32")]

    class Binary1(Module):
        def __init__(self, op):
            super().__init__()
            self.op = op

        def forward(self, lhs, rhs):
            return self.op(lhs, rhs)

    @tvm.script.ir_module
    class expected_binary1:
        @R.function
        def main(
            lhs: R.Tensor((1, 3, 10, 10), dtype="int32"),
            rhs: R.Tensor((1, 3, 10, 10), dtype="int32"),
        ) -> R.Tensor((1, 3, 10, 10), dtype="int32"):
            with R.dataflow():
                lv: R.Tensor((1, 3, 10, 10), dtype="int32") = relax_op(lhs, rhs)
                gv: R.Tensor((1, 3, 10, 10), dtype="int32") = lv
                R.output(gv)
            return gv

    class Binary2(Module):
        def __init__(self, op):
            super().__init__()
            self.op = op

        def forward(self, lhs):
            return self.op(lhs, 1)

    @tvm.script.ir_module
    class expected_binary2:
        @R.function
        def main(
            lhs: R.Tensor((1, 3, 10, 10), dtype="int32"),
        ) -> R.Tensor((1, 3, 10, 10), dtype="int32"):
            with R.dataflow():
                lv: R.Tensor((1, 3, 10, 10), dtype="int32") = relax_op(lhs, R.const(1))
                gv: R.Tensor((1, 3, 10, 10), dtype="int32") = lv
                R.output(gv)
            return gv

    verify_model(Binary1(op), input_info1, {}, expected_binary1)
    verify_model(Binary2(op), input_info2, {}, expected_binary2)


# RSub
def test_rsub():
    input_info1 = [([10, 10], "float32"), ([10, 10], "float32")]
    input_info2 = [([10, 10], "float32")]

    class RSub1(Module):
        def forward(self, x, y):
            return torch.rsub(x, y)

    class RSub2(Module):
        def forward(self, x):
            return torch.rsub(x, 5.0)

    @tvm.script.ir_module
    class expected_rsub1:
        @R.function
        def main(
            x: R.Tensor((10, 10), dtype="float32"), y: R.Tensor((10, 10), dtype="float32")
        ) -> R.Tensor((10, 10), dtype="float32"):
            with R.dataflow():
                lv: R.Tensor((10, 10), dtype="float32") = R.subtract(y, x)
                gv: R.Tensor((10, 10), dtype="float32") = lv
                R.output(gv)
            return gv

    @tvm.script.ir_module
    class expected_rsub2:
        @R.function
        def main(x: R.Tensor((10, 10), dtype="float32")) -> R.Tensor((10, 10), dtype="float32"):
            with R.dataflow():
                lv: R.Tensor((10, 10), dtype="float32") = R.subtract(R.const(5.0, "float32"), x)
                gv: R.Tensor((10, 10), dtype="float32") = lv
                R.output(gv)
            return gv

    verify_model(RSub1(), input_info1, {}, expected_rsub1)
    verify_model(RSub2(), input_info2, {}, expected_rsub2)


# IsIn


def test_isin():
    input_info = [([10, 10], "float32"), ([8], "float32")]

    class IsInModel(torch.nn.Module):
        def forward(self, x, test_elements):
            return torch.isin(x, test_elements)

    @tvm.script.ir_module
    class expected:
        @R.function
        def main(
            inp_0: R.Tensor((10, 10), dtype="float32"), inp_1: R.Tensor((8,), dtype="float32")
        ) -> R.Tensor((10, 10), dtype="bool"):
            with R.dataflow():
                lv: R.Tensor((10, 10, 1), dtype="float32") = R.expand_dims(inp_0, axis=[-1])
                lv1: R.Tensor((8,), dtype="float32") = R.reshape(inp_1, R.shape([8]))
                lv2: R.Tensor((10, 10, 8), dtype="bool") = R.equal(lv, lv1)
                lv3: R.Tensor((10, 10), dtype="bool") = R.sum(lv2, axis=[-1], keepdims=False)
                lv4: R.Tensor((10, 10), dtype="bool") = R.greater(lv3, R.const(0.0, "float32"))
                gv: R.Tensor((10, 10), dtype="bool") = lv4
                R.output(gv)
            return gv

    verify_model(IsInModel(), input_info, {}, expected)


def test_div_mode():
    input_info = [([64, 64], "float32"), ([64, 64], "float32")]

    # Case 1: Basic division (no rounding mode)
    class DivModel(torch.nn.Module):
        def forward(self, x, y):
            return torch.div(x, y)

    @tvm.script.ir_module
    class expected_div:
        @R.function
        def main(
            inp_0: R.Tensor((64, 64), dtype="float32"), inp_1: R.Tensor((64, 64), dtype="float32")
        ) -> R.Tensor((64, 64), dtype="float32"):
            with R.dataflow():
                lv: R.Tensor((64, 64), dtype="float32") = R.divide(inp_0, inp_1)
                gv: R.Tensor((64, 64), dtype="float32") = lv
                R.output(gv)
            return gv

    # Case 2: Division with trunc rounding
    class DivTruncModel(torch.nn.Module):
        def forward(self, x, y):
            return torch.div(x, y, rounding_mode="trunc")

    @tvm.script.ir_module
    class expected_div_trunc:
        @R.function
        def main(
            inp_0: R.Tensor((64, 64), dtype="float32"), inp_1: R.Tensor((64, 64), dtype="float32")
        ) -> R.Tensor((64, 64), dtype="float32"):
            with R.dataflow():
                lv: R.Tensor((64, 64), dtype="float32") = R.divide(inp_0, inp_1)
                lv1: R.Tensor((64, 64), dtype="float32") = R.trunc(lv)
                gv: R.Tensor((64, 64), dtype="float32") = lv1
                R.output(gv)
            return gv

    # Case 3: Division with floor rounding
    class DivFloorModel(torch.nn.Module):
        def forward(self, x, y):
            return torch.div(x, y, rounding_mode="floor")

    @tvm.script.ir_module
    class expected_div_floor:
        @R.function
        def main(
            inp_0: R.Tensor((64, 64), dtype="float32"), inp_1: R.Tensor((64, 64), dtype="float32")
        ) -> R.Tensor((64, 64), dtype="float32"):
            with R.dataflow():
                lv: R.Tensor((64, 64), dtype="float32") = R.floor_divide(inp_0, inp_1)
                gv: R.Tensor((64, 64), dtype="float32") = lv
                R.output(gv)
            return gv

    verify_model(DivModel(), input_info, {}, expected_div)
    verify_model(DivTruncModel(), input_info, {}, expected_div_trunc)
    verify_model(DivFloorModel(), input_info, {}, expected_div_floor)


def test_size():
    input_info = [([1, 3, 10, 10], "float32")]

    class Size(Module):
        def forward(self, input):
            return input.size()

    @tvm.script.ir_module
    class expected1:
        @R.function
        def main(input_1: R.Tensor((1, 3, 10, 10), dtype="float32")) -> R.Shape([1, 3, 10, 10]):
            # block 0
            with R.dataflow():
                gv: R.Shape([1, 3, 10, 10]) = R.shape([1, 3, 10, 10])
                R.output(gv)
            return gv

    verify_model(Size(), input_info, {}, expected1)


def test_squeeze():
    input_info = [([3, 1, 4, 1], "float32")]

    class Squeeze1(Module):
        def forward(self, input):
            return input.squeeze(1)

    @tvm.script.ir_module
    class Expected1:
        @R.function
        def main(inp_0: R.Tensor((3, 1, 4, 1), dtype="float32")) -> R.Tensor(
            (3, 4, 1), dtype="float32"
        ):
            with R.dataflow():
                lv: R.Tensor((3, 4, 1), dtype="float32") = R.squeeze(inp_0, axis=[1])
                gv: R.Tensor((3, 4, 1), dtype="float32") = lv
                R.output(gv)
            return gv

    class Squeeze2(Module):
        def forward(self, input):
            return input.squeeze()

    @tvm.script.ir_module
    class Expected2:
        @R.function
        def main(inp_0: R.Tensor((3, 1, 4, 1), dtype="float32")) -> R.Tensor(
            (3, 4), dtype="float32"
        ):
            with R.dataflow():
                lv: R.Tensor((3, 4), dtype="float32") = R.squeeze(inp_0, axis=None)
                gv: R.Tensor((3, 4), dtype="float32") = lv
                R.output(gv)
            return gv

    verify_model(Squeeze1(), input_info, {}, Expected1)
    verify_model(Squeeze2(), input_info, {}, Expected2)


def test_unsqueeze():
    input_info = [([1, 3, 10, 10], "float32")]

    class Unsqueeze1(Module):
        def forward(self, input):
            return input.unsqueeze(1)

    @tvm.script.ir_module
    class expected1:
        @R.function
        def main(input_1: R.Tensor((1, 3, 10, 10), dtype="float32")) -> R.Tensor(
            (1, 1, 3, 10, 10), dtype="float32"
        ):
            # block 0
            with R.dataflow():
                lv: R.Tensor((1, 1, 3, 10, 10), dtype="float32") = R.expand_dims(input_1, 1)
                gv: R.Tensor((1, 1, 3, 10, 10), dtype="float32") = lv
                R.output(gv)
            return gv

    class Unsqueeze2(Module):
        def forward(self, input):
            return input.unsqueeze(-1)

    @tvm.script.ir_module
    class expected2:
        @R.function
        def main(input_1: R.Tensor((1, 3, 10, 10), dtype="float32")) -> R.Tensor(
            (1, 3, 10, 10, 1), dtype="float32"
        ):
            # block 0
            with R.dataflow():
                lv: R.Tensor((1, 3, 10, 10, 1), dtype="float32") = R.expand_dims(input_1, -1)
                gv: R.Tensor((1, 3, 10, 10, 1), dtype="float32") = lv
                R.output(gv)
            return gv

    verify_model(Unsqueeze1(), input_info, {}, expected1)
    verify_model(Unsqueeze2(), input_info, {}, expected2)


def test_getattr():
    input_info = [([1, 3, 10, 10], "float32")]

    class GetAttr1(Module):
        def forward(self, input):
            return input.shape

    @tvm.script.ir_module
    class expected1:
        @R.function
        def main(input_1: R.Tensor((1, 3, 10, 10), dtype="float32")) -> R.Shape([1, 3, 10, 10]):
            # block 0
            with R.dataflow():
                gv: R.Shape([1, 3, 10, 10]) = R.shape([1, 3, 10, 10])
                R.output(gv)
            return gv

    verify_model(GetAttr1(), input_info, {}, expected1)


def test_getitem():
    class Slice1(Module):
        def forward(self, x):
            return x[0, 1::2, :, :3]

    @tvm.script.ir_module
    class expected1:
        @R.function
        def main(x: R.Tensor((1, 3, 10, 10), dtype="float32")) -> R.Tensor(
            (1, 1, 10, 3), dtype="float32"
        ):
            # block 0
            with R.dataflow():
                lv: R.Tensor((1, 1, 10, 3), dtype="float32") = R.strided_slice(
                    x,
                    axes=[0, 1, 2, 3],
                    begin=[0, 1, 0, 0],
                    end=[1, T.int64(3), T.int64(10), 3],
                    strides=[1, 2, 1, 1],
                )
                lv1: R.Tensor((1, 1, 10, 3), dtype="float32") = R.reshape(lv, (1, 1, 10, 3))
                gv: R.Tensor((1, 1, 10, 3), dtype="float32") = lv1
                R.output(gv)
            return gv

    class Slice2(Module):
        def forward(self, x):
            return x[:, None, None, :, None]

    @I.ir_module
    class expected2:
        @R.function
        def main(inp_0: R.Tensor((8, 16), dtype="float32")) -> R.Tensor(
            (8, 1, 1, 16, 1), dtype="float32"
        ):
            with R.dataflow():
                lv: R.Tensor((8, 16), dtype="float32") = R.strided_slice(
                    inp_0, axes=[0, 1], begin=[0, 0], end=[8, 16], strides=[1, 1]
                )
                lv1: R.Tensor((8, 1, 1, 16, 1), dtype="float32") = R.reshape(
                    lv, R.shape([8, 1, 1, 16, 1])
                )
                gv: R.Tensor((8, 1, 1, 16, 1), dtype="float32") = lv1
                R.output(gv)
            return gv

    verify_model(Slice1(), [([1, 3, 10, 10], "float32")], {}, expected1)
    verify_model(Slice2(), [([8, 16], "float32")], {}, expected2)


operator_basic_unary = [
    (torch.abs, R.abs),
    (torch.acos, R.acos),
    (torch.acosh, R.acosh),
    (torch.asin, R.asin),
    (torch.asinh, R.asinh),
    (torch.atan, R.atan),
    (torch.atanh, R.atanh),
    (torch.bitwise_not, R.bitwise_not),
    (torch.ceil, R.ceil),
    (torch.cos, R.cos),
    (torch.cosh, R.cosh),
    (torch.erf, R.erf),
    (torch.exp, R.exp),
    (torch.floor, R.floor),
    (torch.log, R.log),
    (torch.neg, R.negative),
    (torch.round, R.round),
    (torch.rsqrt, R.rsqrt),
    (torch.sin, R.sin),
    (torch.sinh, R.sinh),
    (torch.sign, R.sign),
    (torch.sqrt, R.sqrt),
    (torch.square, R.square),
    (torch.tan, R.tan),
]


@pytest.mark.parametrize("pytorch_op, relax_op", operator_basic_unary)
def test_basic_unary_ops(pytorch_op, relax_op):
    input_info = [([1, 3, 10, 10], "float32")]

    class Unary(Module):
        def forward(self, input):
            return pytorch_op(input)

    @tvm.script.ir_module
    class expected_unary:
        @R.function
        def main(input_1: R.Tensor((1, 3, 10, 10), dtype="float32")) -> R.Tensor(
            (1, 3, 10, 10), dtype="float32"
        ):
            with R.dataflow():
                lv: R.Tensor((1, 3, 10, 10), dtype="float32") = relax_op(input_1)
                gv: R.Tensor((1, 3, 10, 10), dtype="float32") = lv
                R.output(gv)
            return gv

    verify_model(Unary(), input_info, {}, expected_unary)


def test_sqrt_integer_input_fx():
    input_info = [([1, 4], "int64")]

    class SqrtIntModel(Module):
        def forward(self, input):
            return torch.sqrt(input)

    @tvm.script.ir_module
    class expected:
        @R.function
        def main(input_1: R.Tensor((1, 4), dtype="int64")) -> R.Tensor((1, 4), dtype="float32"):
            with R.dataflow():
                lv: R.Tensor((1, 4), dtype="float32") = R.astype(input_1, dtype="float32")
                lv1: R.Tensor((1, 4), dtype="float32") = R.sqrt(lv)
                gv: R.Tensor((1, 4), dtype="float32") = lv1
                R.output(gv)
            return gv

    verify_model(SqrtIntModel(), input_info, {}, expected)


operator_bool_unary = [
    (torch.isnan, R.isnan),
    (torch.isinf, R.isinf),
    (torch.isfinite, R.isfinite),
]


@pytest.mark.parametrize("pytorch_op, relax_op", operator_bool_unary)
def test_bool_unary_ops(pytorch_op, relax_op):
    input_info = [([1, 3, 10, 10], "float32")]

    class Unary(Module):
        def forward(self, input):
            return pytorch_op(input)

    @tvm.script.ir_module
    class expected_unary:
        @R.function
        def main(input_1: R.Tensor((1, 3, 10, 10), dtype="float32")) -> R.Tensor(
            (1, 3, 10, 10), dtype="bool"
        ):
            with R.dataflow():
                lv: R.Tensor((1, 3, 10, 10), dtype="bool") = relax_op(input_1)
                gv: R.Tensor((1, 3, 10, 10), dtype="bool") = lv
                R.output(gv)
            return gv

    verify_model(Unary(), input_info, {}, expected_unary)


def test_extended_unary_ops():
    input_info = [([1, 3, 10, 10], "float32")]

    # celu
    class Celu1(Module):
        def __init__(self):
            super().__init__()
            self.celu = torch.nn.CELU()

        def forward(self, input):
            return self.celu(input)

    class Celu2(Module):
        def forward(self, input):
            return torch.nn.functional.celu(input)

    # alpha * min(0, exp(x / alpha) - 1) + max(0, x)
    @tvm.script.ir_module
    class expected_celu:
        @R.function
        def main(input_1: R.Tensor((1, 3, 10, 10), dtype="float32")) -> R.Tensor(
            (1, 3, 10, 10), dtype="float32"
        ):
            # block 0
            with R.dataflow():
                lv: R.Tensor((1, 3, 10, 10), dtype="float32") = R.exp(input_1)
                lv_div: R.Tensor((1, 3, 10, 10), dtype="float32") = R.divide(
                    lv, R.const(1.0, "float32")
                )
                lv_sub: R.Tensor((1, 3, 10, 10), dtype="float32") = R.subtract(
                    lv_div, R.const(1.0, "float32")
                )
                lv_min: R.Tensor((1, 3, 10, 10), dtype="float32") = R.minimum(
                    R.const(0.0, "float32"), lv_sub
                )
                lv_scaled: R.Tensor((1, 3, 10, 10), dtype="float32") = R.multiply(
                    R.const(1.0, "float32"), lv_min
                )
                lv_relu_x: R.Tensor((1, 3, 10, 10), dtype="float32") = R.nn.relu(input_1)
                lv_celu: R.Tensor((1, 3, 10, 10), dtype="float32") = R.add(lv_scaled, lv_relu_x)
                gv: R.Tensor((1, 3, 10, 10), dtype="float32") = lv_celu
                R.output(gv)
            return gv

    verify_model(Celu1(), input_info, {}, expected_celu)
    verify_model(Celu2(), input_info, {}, expected_celu)

    # clamp
    class Clamp(Module):
        def forward(self, input):
            return torch.clamp(input, min=0.1, max=0.5)

    @tvm.script.ir_module
    class expected_clamp:
        @R.function
        def main(
            input_1: R.Tensor((1, 3, 10, 10), dtype="float32"),
        ) -> R.Tensor((1, 3, 10, 10), dtype="float32"):
            # block 0
            with R.dataflow():
                lv: R.Tensor((1, 3, 10, 10), dtype="float32") = R.clip(input_1, 0.1, 0.5)
                gv: R.Tensor((1, 3, 10, 10), dtype="float32") = lv
                R.output(gv)
            return gv

    verify_model(Clamp(), input_info, {}, expected_clamp)

    class ClampMinOnly(Module):
        def forward(self, input):
            return torch.clamp(input, min=0.5, max=None)

    @tvm.script.ir_module
    class expected_clamp_min_only:
        @R.function
        def main(
            input_1: R.Tensor((1, 3, 10, 10), dtype="float32"),
        ) -> R.Tensor((1, 3, 10, 10), dtype="float32"):
            # block 0
            with R.dataflow():
                lv: R.Tensor((1, 3, 10, 10), dtype="float32") = R.clip(input_1, 0.5, math.inf)
                gv: R.Tensor((1, 3, 10, 10), dtype="float32") = lv
                R.output(gv)
            return gv

    verify_model(ClampMinOnly(), input_info, {}, expected_clamp_min_only)

    class ClampTensors(Module):
        def forward(self, input):
            return torch.clamp(input, min=input, max=input)

    @tvm.script.ir_module
    class expected_clamp_tensors:
        @R.function
        def main(
            inp_0: R.Tensor((1, 3, 10, 10), dtype="float32"),
        ) -> R.Tensor((1, 3, 10, 10), dtype="float32"):
            # block 0
            with R.dataflow():
                lv: R.Tensor((1, 3, 10, 10), dtype="float32") = R.broadcast_to(
                    inp_0, R.shape([1, 3, 10, 10])
                )
                lv1: R.Tensor((1, 3, 10, 10), dtype="float32") = R.maximum(inp_0, lv)
                lv2: R.Tensor((1, 3, 10, 10), dtype="float32") = R.broadcast_to(
                    inp_0, R.shape([1, 3, 10, 10])
                )
                lv3: R.Tensor((1, 3, 10, 10), dtype="float32") = R.minimum(lv1, lv2)
                lv4: R.Tensor((1, 3, 10, 10), dtype="float32") = R.clip(
                    lv3, R.prim_value(T.float64("-inf")), R.prim_value(T.float64("inf"))
                )
                gv: R.Tensor((1, 3, 10, 10), dtype="float32") = lv4
                R.output(gv)
            return gv

    verify_model(ClampTensors(), input_info, {}, expected_clamp_tensors)

    # dropout
    class Dropout1(Module):
        def __init__(self):
            super().__init__()
            self.dropout = torch.nn.Dropout(0.5)

        def forward(self, input):
            return self.dropout(input)

    class Dropout2(Module):
        def forward(self, input):
            return torch.dropout(input, 0.5, train=True)

    @tvm.script.ir_module
    class expected_dropout:
        @R.function
        def main(input_1: R.Tensor((1, 3, 10, 10), dtype="float32")) -> R.Tensor(
            (1, 3, 10, 10), dtype="float32"
        ):
            # block 0
            with R.dataflow():
                gv: R.Tensor((1, 3, 10, 10), dtype="float32") = input_1
                R.output(gv)
            return gv

    verify_model(Dropout1(), input_info, {}, expected_dropout)
    verify_model(Dropout2(), input_info, {}, expected_dropout)

    # elu
    class Elu(Module):
        def __init__(self):
            super().__init__()
            self.elu = torch.nn.ELU()

        def forward(self, input):
            return self.elu(input)

    class Elu2(Module):
        def forward(self, input):
            return torch.nn.functional.elu(input)

    @tvm.script.ir_module
    class expected_elu:
        @R.function
        def main(input_1: R.Tensor((1, 3, 10, 10), dtype="float32")) -> R.Tensor(
            (1, 3, 10, 10), dtype="float32"
        ):
            # block 0
            with R.dataflow():
                lv_exp: R.Tensor((1, 3, 10, 10), dtype="float32") = R.exp(input_1)
                lv_one_minus_exp: R.Tensor((1, 3, 10, 10), dtype="float32") = R.subtract(
                    R.const(1.0, dtype="float32"), lv_exp
                )
                lv_relu_one_minus_exp: R.Tensor((1, 3, 10, 10), dtype="float32") = R.nn.relu(
                    lv_one_minus_exp
                )
                lv_scaled: R.Tensor((1, 3, 10, 10), dtype="float32") = R.multiply(
                    R.const(-1.0, dtype="float32"), lv_relu_one_minus_exp
                )
                lv_relu_x: R.Tensor((1, 3, 10, 10), dtype="float32") = R.nn.relu(input_1)
                lv_elu: R.Tensor((1, 3, 10, 10), dtype="float32") = R.add(lv_scaled, lv_relu_x)
                gv: R.Tensor((1, 3, 10, 10), dtype="float32") = lv_elu
                R.output(gv)
            return gv

    verify_model(Elu(), input_info, {}, expected_elu)
    verify_model(Elu2(), input_info, {}, expected_elu)

    # gelu
    class Gelu(Module):
        def __init__(self):
            super().__init__()
            self.gelu = torch.nn.GELU()

        def forward(self, input):
            return self.gelu(input)

    class Gelu2(Module):
        def forward(self, input):
            return torch.nn.functional.gelu(input)

    @tvm.script.ir_module
    class expected_gelu:
        @R.function
        def main(input_1: R.Tensor((1, 3, 10, 10), dtype="float32")) -> R.Tensor(
            (1, 3, 10, 10), dtype="float32"
        ):
            # block 0
            with R.dataflow():
                lv: R.Tensor((1, 3, 10, 10), dtype="float32") = R.nn.gelu(input_1)
                gv: R.Tensor((1, 3, 10, 10), dtype="float32") = lv
                R.output(gv)
            return gv

    verify_model(Gelu(), input_info, {}, expected_gelu)
    verify_model(Gelu2(), input_info, {}, expected_gelu)

    # hardsigmoid
    class Hardsigmoid(torch.nn.Module):
        def __init__(self):
            super().__init__()
            self.hs = torch.nn.Hardsigmoid()

        def forward(self, input):
            return self.hs(input)

    class Hardsigmoid2(torch.nn.Module):
        def forward(self, input):
            return torch.nn.functional.hardsigmoid(input)

    @tvm.script.ir_module
    class expected_hardsigmoid:
        @R.function
        def main(inp_0: R.Tensor((1, 3, 10, 10), dtype="float32")) -> R.Tensor(
            (1, 3, 10, 10), dtype="float32"
        ):
            with R.dataflow():
                lv: R.Tensor((1, 3, 10, 10), dtype="float32") = R.add(inp_0, R.const(3, "float32"))
                lv1: R.Tensor((1, 3, 10, 10), dtype="float32") = R.clip(lv, 0, 6)
                lv2: R.Tensor((1, 3, 10, 10), dtype="float32") = R.divide(
                    lv1, R.const(6, "float32")
                )
                gv: R.Tensor((1, 3, 10, 10), dtype="float32") = lv2
                R.output(gv)
            return gv

    verify_model(Hardsigmoid(), input_info, {}, expected_hardsigmoid)
    verify_model(Hardsigmoid2(), input_info, {}, expected_hardsigmoid)

    # hardswish
    class Hardswish(torch.nn.Module):
        def __init__(self):
            super().__init__()
            self.hs = torch.nn.Hardswish()

        def forward(self, input):
            return self.hs(input)

    class Hardswish2(torch.nn.Module):
        def forward(self, input):
            return torch.nn.functional.hardswish(input)

    @tvm.script.ir_module
    class expected_hardswish:
        @R.function
        def main(inp_0: R.Tensor((1, 3, 10, 10), dtype="float32")) -> R.Tensor(
            (1, 3, 10, 10), dtype="float32"
        ):
            with R.dataflow():
                lv: R.Tensor((1, 3, 10, 10), dtype="float32") = R.add(inp_0, R.const(3, "float32"))
                lv1: R.Tensor((1, 3, 10, 10), dtype="float32") = R.clip(lv, 0, 6)
                lv2: R.Tensor((1, 3, 10, 10), dtype="float32") = R.divide(
                    lv1, R.const(6, "float32")
                )
                lv3: R.Tensor((1, 3, 10, 10), dtype="float32") = R.multiply(inp_0, lv2)
                gv: R.Tensor((1, 3, 10, 10), dtype="float32") = lv3
                R.output(gv)
            return gv

    verify_model(Hardswish(), input_info, {}, expected_hardswish)
    verify_model(Hardswish2(), input_info, {}, expected_hardswish)

    # hardtanh
    class Hardtanh(torch.nn.Module):
        def __init__(self):
            super().__init__()
            self.ht = torch.nn.Hardtanh()

        def forward(self, input):
            return self.ht(input)

    class Hardtanh2(torch.nn.Module):
        def forward(self, input):
            return torch.nn.functional.hardtanh(input)

    @tvm.script.ir_module
    class expected1:
        @R.function
        def main(inp_0: R.Tensor((1, 3, 10, 10), dtype="float32")) -> R.Tensor(
            (1, 3, 10, 10), dtype="float32"
        ):
            with R.dataflow():
                lv: R.Tensor((1, 3, 10, 10), dtype="float32") = R.clip(
                    inp_0, R.prim_value(T.float64(-1.0)), R.prim_value(T.float64(1.0))
                )
                gv: R.Tensor((1, 3, 10, 10), dtype="float32") = lv
                R.output(gv)
            return gv

    verify_model(Hardtanh(), input_info, {}, expected1)
    verify_model(Hardtanh2(), input_info, {}, expected1)

    # leaky_relu
    test_leakyrelu()

    # softplus
    test_softplus()

    # prelu
    test_prelu()

    # log2
    class Log2(Module):
        def forward(self, x):
            return torch.log2(x)

    @tvm.script.ir_module
    class Expected_log2:
        @R.function
        def main(
            inp_0: R.Tensor((1, 3, 10, 10), dtype="float32"),
        ) -> R.Tensor((1, 3, 10, 10), dtype="float32"):
            with R.dataflow():
                lv: R.Tensor((1, 3, 10, 10), dtype="float32") = R.log(inp_0)
                lv1: R.Tensor((1, 3, 10, 10), dtype="float32") = R.divide(
                    lv, R.const(0.6931471805599453, "float32")
                )
                gv: R.Tensor((1, 3, 10, 10), dtype="float32") = lv1
                R.output(gv)
            return gv

    verify_model(Log2(), input_info, {}, Expected_log2)

    # log10
    class Log10(Module):
        def forward(self, x):
            return torch.log10(x)

    @tvm.script.ir_module
    class Expected_log10:
        @R.function
        def main(
            inp_0: R.Tensor((1, 3, 10, 10), dtype="float32"),
        ) -> R.Tensor((1, 3, 10, 10), dtype="float32"):
            with R.dataflow():
                lv: R.Tensor((1, 3, 10, 10), dtype="float32") = R.log(inp_0)
                lv1: R.Tensor((1, 3, 10, 10), dtype="float32") = R.divide(
                    lv, R.const(2.302585092994046, "float32")
                )
                gv: R.Tensor((1, 3, 10, 10), dtype="float32") = lv1
                R.output(gv)
            return gv

    verify_model(Log10(), input_info, {}, Expected_log10)

    # log1p
    class Log1p(Module):
        def forward(self, x):
            return torch.log1p(x)

    @tvm.script.ir_module
    class Expected_log1p:
        @R.function
        def main(
            inp_0: R.Tensor((1, 3, 10, 10), dtype="float32"),
        ) -> R.Tensor((1, 3, 10, 10), dtype="float32"):
            with R.dataflow():
                lv: R.Tensor((1, 3, 10, 10), dtype="float32") = R.log(
                    R.add(inp_0, R.const(1, "float32"))
                )
                gv: R.Tensor((1, 3, 10, 10), dtype="float32") = lv
                R.output(gv)
            return gv

    verify_model(Log1p(), input_info, {}, Expected_log1p)

    # logical_not
    class LogicalNot(Module):
        def forward(self, input):
            return torch.logical_not(input)

    @tvm.script.ir_module
    class expected_logical_not:
        @R.function
        def main(inp_0: R.Tensor((1, 3, 10, 10), dtype="float32")) -> R.Tensor(
            (1, 3, 10, 10), dtype="bool"
        ):
            with R.dataflow():
                lv: R.Tensor((1, 3, 10, 10), dtype="bool") = R.astype(inp_0, dtype="bool")
                lv1: R.Tensor((1, 3, 10, 10), dtype="bool") = R.logical_not(lv)
                gv: R.Tensor((1, 3, 10, 10), dtype="bool") = lv1
                R.output(gv)
            return gv

    verify_model(LogicalNot(), input_info, {}, expected_logical_not)

    # log_softmax
    class LogSoftmax(Module):
        def __init__(self):
            super().__init__()
            self.lsm = torch.nn.LogSoftmax(dim=1)

        def forward(self, input):
            return self.lsm(input)

    class LogSoftmax2(Module):
        def forward(self, input):
            return torch.nn.functional.log_softmax(input, dim=1)

    @tvm.script.ir_module
    class expected_log_softmax:
        @R.function
        def main(input_1: R.Tensor((1, 3, 10, 10), dtype="float32")) -> R.Tensor(
            (1, 3, 10, 10), dtype="float32"
        ):
            # block 0
            with R.dataflow():
                lv: R.Tensor((1, 3, 10, 10), dtype="float32") = R.nn.log_softmax(input_1, axis=1)
                gv: R.Tensor((1, 3, 10, 10), dtype="float32") = lv
                R.output(gv)
            return gv

    verify_model(LogSoftmax(), input_info, {}, expected_log_softmax)
    verify_model(LogSoftmax2(), input_info, {}, expected_log_softmax)

    # reciprocal
    class Reciprocal(Module):
        def forward(self, input):
            return torch.reciprocal(input)

    @tvm.script.ir_module
    class expected_reciprocal:
        @R.function
        def main(input_1: R.Tensor((1, 3, 10, 10), dtype="float32")) -> R.Tensor(
            (1, 3, 10, 10), dtype="float32"
        ):
            with R.dataflow():
                lv: R.Tensor((1, 3, 10, 10), dtype="float32") = R.divide(
                    R.const(1.0, "float32"), input_1
                )
                gv: R.Tensor((1, 3, 10, 10), dtype="float32") = lv
                R.output(gv)
            return gv

    verify_model(Reciprocal(), input_info, {}, expected_reciprocal)

    # relu
    class ReLU0(Module):
        def __init__(self):
            super().__init__()
            self.relu = torch.nn.ReLU()

        def forward(self, input):
            return self.relu(input)

    class ReLU1(Module):
        def forward(self, input):
            return torch.nn.functional.relu(input)

    @tvm.script.ir_module
    class expected_relu:
        @R.function
        def main(input_1: R.Tensor((1, 3, 10, 10), dtype="float32")) -> R.Tensor(
            (1, 3, 10, 10), dtype="float32"
        ):
            # block 0
            with R.dataflow():
                lv: R.Tensor((1, 3, 10, 10), dtype="float32") = R.nn.relu(input_1)
                gv: R.Tensor((1, 3, 10, 10), dtype="float32") = lv
                R.output(gv)
            return gv

    verify_model(ReLU0(), input_info, {}, expected_relu)
    verify_model(ReLU1(), input_info, {}, expected_relu)

    # relu6
    class ReLU6_1(torch.nn.Module):
        def __init__(self):
            super().__init__()
            self.relu6 = torch.nn.ReLU6()

        def forward(self, x):
            return self.relu6(x)

    class ReLU6_2(torch.nn.Module):
        def forward(self, x):
            return torch.nn.functional.relu6(x)

    @tvm.script.ir_module
    class expected:
        @R.function
        def main(inp_0: R.Tensor((1, 3, 10, 10), dtype="float32")) -> R.Tensor(
            (1, 3, 10, 10), dtype="float32"
        ):
            with R.dataflow():
                lv: R.Tensor((1, 3, 10, 10), dtype="float32") = R.nn.relu6(inp_0)
                gv: R.Tensor((1, 3, 10, 10), dtype="float32") = lv
                R.output(gv)
            return gv

    verify_model(ReLU6_1(), input_info, {}, expected)
    verify_model(ReLU6_2(), input_info, {}, expected)

    # selu
    class Selu1(Module):
        def __init__(self):
            super().__init__()
            self.selu = torch.nn.SELU()

        def forward(self, input):
            return self.selu(input)

    class Selu2(Module):
        def forward(self, input):
            return torch.nn.functional.selu(input)

    @tvm.script.ir_module
    class expected_selu:
        @R.function
        def main(inp_0: R.Tensor((1, 3, 10, 10), dtype="float32")) -> R.Tensor(
            (1, 3, 10, 10), dtype="float32"
        ):
            with R.dataflow():
                lv: R.Tensor((1, 3, 10, 10), dtype="float32") = R.nn.selu(inp_0)
                gv: R.Tensor((1, 3, 10, 10), dtype="float32") = lv
                R.output(gv)
            return gv

    verify_model(Selu1(), input_info, {}, expected_selu)
    verify_model(Selu2(), input_info, {}, expected_selu)

    # sigmoid
    class Sigmoid(Module):
        def __init__(self):
            super().__init__()
            self.sigmoid = torch.nn.Sigmoid()

        def forward(self, input):
            return self.sigmoid(input)

    class Sigmoid2(Module):
        def forward(self, input):
            return torch.sigmoid(input)

    @tvm.script.ir_module
    class expected_sigmoid:
        @R.function
        def main(input_1: R.Tensor((1, 3, 10, 10), dtype="float32")) -> R.Tensor(
            (1, 3, 10, 10), dtype="float32"
        ):
            # block 0
            with R.dataflow():
                lv: R.Tensor((1, 3, 10, 10), dtype="float32") = R.sigmoid(input_1)
                gv: R.Tensor((1, 3, 10, 10), dtype="float32") = lv
                R.output(gv)
            return gv

    verify_model(Sigmoid(), input_info, {}, expected_sigmoid)
    verify_model(Sigmoid2(), input_info, {}, expected_sigmoid)

    # silu
    class SiLU(Module):
        def __init__(self):
            super().__init__()
            self.silu = torch.nn.SiLU()

        def forward(self, input):
            return self.silu(input)

    class SiLU2(Module):
        def forward(self, input):
            return torch.nn.functional.silu(input)

    @tvm.script.ir_module
    class expected_silu:
        @R.function
        def main(input_1: R.Tensor((1, 3, 10, 10), dtype="float32")) -> R.Tensor(
            (1, 3, 10, 10), dtype="float32"
        ):
            # block 0
            with R.dataflow():
                lv: R.Tensor((1, 3, 10, 10), dtype="float32") = R.nn.silu(input_1)
                gv: R.Tensor((1, 3, 10, 10), dtype="float32") = lv
                R.output(gv)
            return gv

    verify_model(SiLU(), input_info, {}, expected_silu)
    verify_model(SiLU2(), input_info, {}, expected_silu)

    # softmax
    class Softmax(Module):
        def __init__(self):
            super().__init__()
            self.sm = torch.nn.Softmax(dim=1)

        def forward(self, input):
            return self.sm(input)

    class Softmax2(Module):
        def forward(self, input):
            return torch.nn.functional.softmax(input, dim=1)

    @tvm.script.ir_module
    class expected_softmax:
        @R.function
        def main(input_1: R.Tensor((1, 3, 10, 10), dtype="float32")) -> R.Tensor(
            (1, 3, 10, 10), dtype="float32"
        ):
            # block 0
            with R.dataflow():
                lv: R.Tensor((1, 3, 10, 10), dtype="float32") = R.nn.softmax(input_1, axis=1)
                gv: R.Tensor((1, 3, 10, 10), dtype="float32") = lv
                R.output(gv)
            return gv

    verify_model(Softmax(), input_info, {}, expected_softmax)
    verify_model(Softmax2(), input_info, {}, expected_softmax)

    # tanh
    class Tanh(Module):
        def __init__(self):
            super().__init__()
            self.tanh = torch.nn.Tanh()

        def forward(self, input):
            return self.tanh(input)

    class Tanh2(Module):
        def forward(self, input):
            return torch.tanh(input)

    @tvm.script.ir_module
    class expected_tanh:
        @R.function
        def main(input_1: R.Tensor((1, 3, 10, 10), dtype="float32")) -> R.Tensor(
            (1, 3, 10, 10), dtype="float32"
        ):
            # block 0
            with R.dataflow():
                lv: R.Tensor((1, 3, 10, 10), dtype="float32") = R.tanh(input_1)
                gv: R.Tensor((1, 3, 10, 10), dtype="float32") = lv
                R.output(gv)
            return gv

    verify_model(Tanh(), input_info, {}, expected_tanh)
    verify_model(Tanh2(), input_info, {}, expected_tanh)

    # tril
    class Tril(Module):
        def forward(self, input):
            return torch.tril(input, 1)

    class InplaceTril(Module):
        def forward(self, input):
            input.tril_(1)
            return input

    @tvm.script.ir_module
    class expected_tril:
        @R.function
        def main(input_1: R.Tensor((1, 3, 10, 10), dtype="float32")) -> R.Tensor(
            (1, 3, 10, 10), dtype="float32"
        ):
            with R.dataflow():
                lv: R.Tensor((1, 3, 10, 10), dtype="float32") = R.tril(input_1, 1)
                gv: R.Tensor((1, 3, 10, 10), dtype="float32") = lv
                R.output(gv)
            return gv

    verify_model(Tril(), input_info, {}, expected_tril)
    verify_model(InplaceTril(), input_info, {}, expected_tril)

    # triu
    class Triu(Module):
        def forward(self, input):
            return torch.triu(input, 1)

    class InplaceTriu(Module):
        def forward(self, input):
            input.triu_(1)
            return input

    @tvm.script.ir_module
    class expected_triu:
        @R.function
        def main(input_1: R.Tensor((1, 3, 10, 10), dtype="float32")) -> R.Tensor(
            (1, 3, 10, 10), dtype="float32"
        ):
            with R.dataflow():
                lv: R.Tensor((1, 3, 10, 10), dtype="float32") = R.triu(input_1, 1)
                gv: R.Tensor((1, 3, 10, 10), dtype="float32") = lv
                R.output(gv)
            return gv

    verify_model(Triu(), input_info, {}, expected_triu)
    verify_model(InplaceTriu(), input_info, {}, expected_triu)

    # trunc
    class Trunc(torch.nn.Module):
        def forward(self, input):
            return torch.trunc(input)

    @tvm.script.ir_module
    class expected_trunc:
        @R.function
        def main(inp_0: R.Tensor((1, 3, 10, 10), dtype="float32")) -> R.Tensor(
            (1, 3, 10, 10), dtype="float32"
        ):
            with R.dataflow():
                lv: R.Tensor((1, 3, 10, 10), dtype="float32") = R.trunc(inp_0)
                gv: R.Tensor((1, 3, 10, 10), dtype="float32") = lv
                R.output(gv)
            return gv

    verify_model(Trunc(), input_info, {}, expected_trunc)


def test_logical_and():
    input_info = [([1, 3, 10, 10], "float32"), ([1, 3, 10, 10], "float32")]

    class LogicalAnd(Module):
        def forward(self, lhs, rhs):
            return torch.logical_and(lhs, rhs)

    @tvm.script.ir_module
    class expected:
        @R.function
        def main(
            lhs: R.Tensor((1, 3, 10, 10), dtype="float32"),
            rhs: R.Tensor((1, 3, 10, 10), dtype="float32"),
        ) -> R.Tensor((1, 3, 10, 10), dtype="bool"):
            with R.dataflow():
                lv: R.Tensor((1, 3, 10, 10), dtype="bool") = R.astype(lhs, dtype="bool")
                lv1: R.Tensor((1, 3, 10, 10), dtype="bool") = R.astype(rhs, dtype="bool")
                lv2: R.Tensor((1, 3, 10, 10), dtype="bool") = R.logical_and(lv, lv1)
                gv: R.Tensor((1, 3, 10, 10), dtype="bool") = lv2
                R.output(gv)
            return gv

    verify_model(LogicalAnd(), input_info, {}, expected)


def test_logical_or():
    input_info = [([1, 3, 10, 10], "float32"), ([1, 3, 10, 10], "float32")]

    class LogicalOr(Module):
        def forward(self, lhs, rhs):
            return torch.logical_or(lhs, rhs)

    @tvm.script.ir_module
    class expected:
        @R.function
        def main(
            lhs: R.Tensor((1, 3, 10, 10), dtype="float32"),
            rhs: R.Tensor((1, 3, 10, 10), dtype="float32"),
        ) -> R.Tensor((1, 3, 10, 10), dtype="bool"):
            with R.dataflow():
                lv: R.Tensor((1, 3, 10, 10), dtype="bool") = R.astype(lhs, dtype="bool")
                lv1: R.Tensor((1, 3, 10, 10), dtype="bool") = R.astype(rhs, dtype="bool")
                lv2: R.Tensor((1, 3, 10, 10), dtype="bool") = R.logical_or(lv, lv1)
                gv: R.Tensor((1, 3, 10, 10), dtype="bool") = lv2
                R.output(gv)
            return gv

    verify_model(LogicalOr(), input_info, {}, expected)


def test_logical_xor():
    input_info = [([1, 3, 10, 10], "float32"), ([1, 3, 10, 10], "float32")]

    class LogicalXor(Module):
        def forward(self, lhs, rhs):
            return torch.logical_xor(lhs, rhs)

    @tvm.script.ir_module
    class expected:
        @R.function
        def main(
            lhs: R.Tensor((1, 3, 10, 10), dtype="float32"),
            rhs: R.Tensor((1, 3, 10, 10), dtype="float32"),
        ) -> R.Tensor((1, 3, 10, 10), dtype="bool"):
            with R.dataflow():
                lv: R.Tensor((1, 3, 10, 10), dtype="bool") = R.astype(lhs, dtype="bool")
                lv1: R.Tensor((1, 3, 10, 10), dtype="bool") = R.astype(rhs, dtype="bool")
                lv2: R.Tensor((1, 3, 10, 10), dtype="bool") = R.logical_xor(lv, lv1)
                gv: R.Tensor((1, 3, 10, 10), dtype="bool") = lv2
                R.output(gv)
            return gv

    verify_model(LogicalXor(), input_info, {}, expected)


def test_pow_integer():
    input_info = [([4], "int64")]

    class Pow(Module):
        def forward(self, input):
            return input.pow(4)

    @tvm.script.ir_module
    class expected:
        @R.function
        def main(inp_0: R.Tensor((4,), dtype="int64")) -> R.Tensor((4,), dtype="int64"):
            with R.dataflow():
                lv: R.Tensor((4,), dtype="int64") = R.multiply(inp_0, inp_0)
                lv1: R.Tensor((4,), dtype="int64") = R.multiply(lv, inp_0)
                lv2: R.Tensor((4,), dtype="int64") = R.multiply(lv1, inp_0)
                gv: R.Tensor((4,), dtype="int64") = lv2
                R.output(gv)
            return gv

    verify_model(Pow(), input_info, {}, expected)


def test_interpolate():
    input_info = [([1, 3, 10, 10], "float32")]

    class Interpolate(Module):
        def forward(self, input):
            return torch.nn.functional.interpolate(input, (5, 5))

    @tvm.script.ir_module
    class expected1:
        @R.function
        def main(input_1: R.Tensor((1, 3, 10, 10), dtype="float32")) -> R.Tensor(
            (1, 3, 5, 5), dtype="float32"
        ):
            # block 0
            with R.dataflow():
                lv: R.Tensor((1, 3, 5, 5), dtype="float32") = R.image.resize2d(
                    input_1,
                    (5, 5),
                    roi=[0.000000, 0.000000, 0.000000, 0.000000],
                    layout="NCHW",
                    method="nearest_neighbor",
                    coordinate_transformation_mode="asymmetric",
                    rounding_method="round",
                    cubic_alpha=-0.75,
                    cubic_exclude=0,
                    extrapolation_value=0,
                )
                gv: R.Tensor((1, 3, 5, 5), dtype="float32") = lv
                R.output(gv)
            return gv

    verify_model(Interpolate(), input_info, {}, expected1)

    class Interpolate2(Module):
        def forward(self, input):
            return torch.nn.functional.interpolate(
                input,
                size=None,
                scale_factor=2.0,
                mode="bilinear",
                align_corners=False,
            )

    @tvm.script.ir_module
    class expected2:
        @R.function
        def main(input_1: R.Tensor((1, 3, 10, 10), dtype="float32")) -> R.Tensor(
            (1, 3, 20, 20), dtype="float32"
        ):
            # block 0
            with R.dataflow():
                lv: R.Tensor((1, 3, 20, 20), dtype="float32") = R.image.resize2d(
                    input_1,
                    (20, 20),
                    roi=[0.000000, 0.000000, 0.000000, 0.000000],
                    layout="NCHW",
                    method="linear",
                    coordinate_transformation_mode="half_pixel",
                    rounding_method="round",
                    cubic_alpha=-0.75,
                    cubic_exclude=0,
                    extrapolation_value=0,
                )
                gv: R.Tensor((1, 3, 20, 20), dtype="float32") = lv
                R.output(gv)
            return gv

    verify_model(Interpolate2(), input_info, {}, expected2)

    class Interpolate3(Module):
        def forward(self, input):
            return torch.nn.functional.interpolate(
                input,
                size=None,
                scale_factor=(2.0, 1.0),
                mode="bilinear",
                align_corners=False,
            )

    @tvm.script.ir_module
    class expected3:
        @R.function
        def main(input_1: R.Tensor((1, 3, 10, 10), dtype="float32")) -> R.Tensor(
            (1, 3, 20, 10), dtype="float32"
        ):
            # block 0
            with R.dataflow():
                lv: R.Tensor((1, 3, 20, 10), dtype="float32") = R.image.resize2d(
                    input_1,
                    (20, 10),
                    roi=[0.000000, 0.000000, 0.000000, 0.000000],
                    layout="NCHW",
                    method="linear",
                    coordinate_transformation_mode="half_pixel",
                    rounding_method="round",
                    cubic_alpha=-0.75,
                    cubic_exclude=0,
                    extrapolation_value=0,
                )
                gv: R.Tensor((1, 3, 20, 10), dtype="float32") = lv
                R.output(gv)
            return gv

    verify_model(Interpolate3(), input_info, {}, expected3)

    class Interpolate4(Module):
        def forward(self, input):
            return torch.nn.functional.interpolate(
                input,
                size=None,
                scale_factor=(2.0, 1.0),
                mode="bicubic",
                align_corners=False,
            )

    @tvm.script.ir_module
    class expected4:
        @R.function
        def main(input_1: R.Tensor((1, 3, 10, 10), dtype="float32")) -> R.Tensor(
            (1, 3, 20, 10), dtype="float32"
        ):
            # block 0
            with R.dataflow():
                lv: R.Tensor((1, 3, 20, 10), dtype="float32") = R.image.resize2d(
                    input_1,
                    (20, 10),
                    roi=[0.000000, 0.000000, 0.000000, 0.000000],
                    layout="NCHW",
                    method="cubic",
                    coordinate_transformation_mode="half_pixel",
                    rounding_method="round",
                    cubic_alpha=-0.75,
                    cubic_exclude=0,
                    extrapolation_value=0,
                )
                gv: R.Tensor((1, 3, 20, 10), dtype="float32") = lv
                R.output(gv)
            return gv

    verify_model(Interpolate4(), input_info, {}, expected4)

    input_info_5d = [([1, 3, 4, 10, 10], "float32")]

    class Interpolate5(Module):
        def forward(self, input):
            return torch.nn.functional.interpolate(
                input,
                size=None,
                scale_factor=(2.0, 2.0, 2.0),
                mode="trilinear",
                align_corners=False,
            )

    @tvm.script.ir_module
    class expected5:
        @R.function
        def main(input_5: R.Tensor((1, 3, 4, 10, 10), dtype="float32")) -> R.Tensor(
            (1, 3, 8, 20, 20), dtype="float32"
        ):
            with R.dataflow():
                lv: R.Tensor((1, 3, 8, 20, 20), dtype="float32") = R.image.resize3d(
                    input_5,
                    (8, 20, 20),
                    roi=[0.000000, 0.000000, 0.000000, 0.000000, 0.000000, 0.000000],
                    layout="NCDHW",
                    method="linear",
                    coordinate_transformation_mode="half_pixel",
                    rounding_method="",
                    cubic_alpha=-0.75,
                    cubic_exclude=0,
                    extrapolation_value=0,
                )
                gv: R.Tensor((1, 3, 8, 20, 20), dtype="float32") = lv
                R.output(gv)
            return gv

    verify_model(Interpolate5(), input_info_5d, {}, expected5)

    class Interpolate6(Module):
        def forward(self, input):
            return torch.nn.functional.interpolate(
                input,
                size=None,
                scale_factor=(2.0, 4.0, 4.0),
                mode="trilinear",
                align_corners=False,
            )

    @tvm.script.ir_module
    class expected6:
        @R.function
        def main(input_5: R.Tensor((1, 3, 4, 10, 10), dtype="float32")) -> R.Tensor(
            (1, 3, 8, 40, 40), dtype="float32"
        ):
            with R.dataflow():
                lv: R.Tensor((1, 3, 8, 40, 40), dtype="float32") = R.image.resize3d(
                    input_5,
                    (8, 40, 40),
                    roi=[0.000000, 0.000000, 0.000000, 0.000000, 0.000000, 0.000000],
                    layout="NCDHW",
                    method="linear",
                    coordinate_transformation_mode="half_pixel",
                    rounding_method="",
                    cubic_alpha=-0.75,
                    cubic_exclude=0,
                    extrapolation_value=0,
                )
                gv: R.Tensor((1, 3, 8, 40, 40), dtype="float32") = lv
                R.output(gv)
            return gv

    verify_model(Interpolate6(), input_info_5d, {}, expected6)

    class Interpolate7(Module):
        def forward(self, input):
            return torch.nn.functional.interpolate(
                input,
                size=(8, 40, 40),
                mode="trilinear",
                align_corners=False,
            )

    @tvm.script.ir_module
    class expected7:
        @R.function
        def main(input_5: R.Tensor((1, 3, 4, 10, 10), dtype="float32")) -> R.Tensor(
            (1, 3, 8, 40, 40), dtype="float32"
        ):
            with R.dataflow():
                lv: R.Tensor((1, 3, 8, 40, 40), dtype="float32") = R.image.resize3d(
                    input_5,
                    (8, 40, 40),
                    roi=[0.000000, 0.000000, 0.000000, 0.000000, 0.000000, 0.000000],
                    layout="NCDHW",
                    method="linear",
                    coordinate_transformation_mode="half_pixel",
                    rounding_method="",
                    cubic_alpha=-0.75,
                    cubic_exclude=0,
                    extrapolation_value=0,
                )
                gv: R.Tensor((1, 3, 8, 40, 40), dtype="float32") = lv
                R.output(gv)
            return gv

    verify_model(Interpolate7(), input_info_5d, {}, expected7)

    class Interpolate8(Module):
        def forward(self, input):
            return torch.nn.functional.interpolate(
                input,
                size=(8, 40, 40),
                mode="trilinear",
                align_corners=True,
            )

    @tvm.script.ir_module
    class expected8:
        @R.function
        def main(input_5: R.Tensor((1, 3, 4, 10, 10), dtype="float32")) -> R.Tensor(
            (1, 3, 8, 40, 40), dtype="float32"
        ):
            with R.dataflow():
                lv: R.Tensor((1, 3, 8, 40, 40), dtype="float32") = R.image.resize3d(
                    input_5,
                    (8, 40, 40),
                    roi=[0.000000, 0.000000, 0.000000, 0.000000, 0.000000, 0.000000],
                    layout="NCDHW",
                    method="linear",
                    coordinate_transformation_mode="align_corners",
                    rounding_method="",
                    cubic_alpha=-0.75,
                    cubic_exclude=0,
                    extrapolation_value=0,
                )
                gv: R.Tensor((1, 3, 8, 40, 40), dtype="float32") = lv
                R.output(gv)
            return gv

    verify_model(Interpolate8(), input_info_5d, {}, expected8)


def test_interpolate_nhwc_layout():
    # First verify backward compatibility - default should still be NCHW
    input_info_nchw = [([1, 3, 10, 10], "float32")]

    class InterpolateDefault(Module):
        def forward(self, input):
            return torch.nn.functional.interpolate(input, (5, 5))

    @tvm.script.ir_module
    class expected_default_nchw:
        @R.function
        def main(input_1: R.Tensor((1, 3, 10, 10), dtype="float32")) -> R.Tensor(
            (1, 3, 5, 5), dtype="float32"
        ):
            # block 0
            with R.dataflow():
                lv: R.Tensor((1, 3, 5, 5), dtype="float32") = R.image.resize2d(
                    input_1,
                    (5, 5),
                    roi=[0.000000, 0.000000, 0.000000, 0.000000],
                    layout="NCHW",
                    method="nearest_neighbor",
                    coordinate_transformation_mode="asymmetric",
                    rounding_method="round",
                    cubic_alpha=-0.75,
                    cubic_exclude=0,
                    extrapolation_value=0,
                )
                gv: R.Tensor((1, 3, 5, 5), dtype="float32") = lv
                R.output(gv)
            return gv

    # Verify default behavior (no default_image_layout parameter) uses NCHW
    graph_model_default = fx.symbolic_trace(InterpolateDefault())
    with torch.no_grad():
        mod_default = from_fx(graph_model_default, input_info_nchw)
    tvm.ir.assert_structural_equal(mod_default, expected_default_nchw)

    # Now test NHWC layout
    input_info = [([1, 10, 10, 3], "float32")]

    class InterpolateNHWC(Module):
        def forward(self, input):
            return torch.nn.functional.interpolate(input, (5, 5))

    @tvm.script.ir_module
    class expected_nhwc:
        @R.function
        def main(input_1: R.Tensor((1, 10, 10, 3), dtype="float32")) -> R.Tensor(
            (1, 5, 5, 3), dtype="float32"
        ):
            # block 0
            with R.dataflow():
                lv: R.Tensor((1, 5, 5, 3), dtype="float32") = R.image.resize2d(
                    input_1,
                    (5, 5),
                    roi=[0.000000, 0.000000, 0.000000, 0.000000],
                    layout="NHWC",
                    method="nearest_neighbor",
                    coordinate_transformation_mode="asymmetric",
                    rounding_method="round",
                    cubic_alpha=-0.75,
                    cubic_exclude=0,
                    extrapolation_value=0,
                )
                gv: R.Tensor((1, 5, 5, 3), dtype="float32") = lv
                R.output(gv)
            return gv

    # Test with NHWC layout
    graph_model = fx.symbolic_trace(InterpolateNHWC())
    with torch.no_grad():
        mod = from_fx(graph_model, input_info, default_image_layout="NHWC")
    tvm.ir.assert_structural_equal(mod, expected_nhwc)

    # Test with bilinear interpolation and NHWC layout
    class InterpolateNHWC2(Module):
        def forward(self, input):
            return torch.nn.functional.interpolate(
                input, size=None, scale_factor=2.0, mode="bilinear", align_corners=False
            )

    @tvm.script.ir_module
    class expected_nhwc2:
        @R.function
        def main(input_1: R.Tensor((1, 10, 10, 3), dtype="float32")) -> R.Tensor(
            (1, 20, 20, 3), dtype="float32"
        ):
            # block 0
            with R.dataflow():
                lv: R.Tensor((1, 20, 20, 3), dtype="float32") = R.image.resize2d(
                    input_1,
                    (20, 20),
                    roi=[0.000000, 0.000000, 0.000000, 0.000000],
                    layout="NHWC",
                    method="linear",
                    coordinate_transformation_mode="half_pixel",
                    rounding_method="round",
                    cubic_alpha=-0.75,
                    cubic_exclude=0,
                    extrapolation_value=0,
                )
                gv: R.Tensor((1, 20, 20, 3), dtype="float32") = lv
                R.output(gv)
            return gv

    graph_model2 = fx.symbolic_trace(InterpolateNHWC2())
    with torch.no_grad():
        mod2 = from_fx(graph_model2, input_info, default_image_layout="NHWC")
    tvm.ir.assert_structural_equal(mod2, expected_nhwc2)

    input_info_5d = [([1, 4, 10, 10, 3], "float32")]

    class InterpolateNHWC3(Module):
        def forward(self, input):
            return torch.nn.functional.interpolate(
                input,
                size=None,
                scale_factor=(2.0, 4.0, 4.0),
                mode="trilinear",
                align_corners=False,
            )

    @tvm.script.ir_module
    class expected_nhwc3:
        @R.function
        def main(input_5: R.Tensor((1, 4, 10, 10, 3), dtype="float32")) -> R.Tensor(
            (1, 8, 40, 40, 3), dtype="float32"
        ):
            with R.dataflow():
                lv: R.Tensor((1, 8, 40, 40, 3), dtype="float32") = R.image.resize3d(
                    input_5,
                    (8, 40, 40),
                    roi=[0.000000, 0.000000, 0.000000, 0.000000, 0.000000, 0.000000],
                    layout="NDHWC",
                    method="linear",
                    coordinate_transformation_mode="half_pixel",
                    rounding_method="",
                    cubic_alpha=-0.75,
                    cubic_exclude=0,
                    extrapolation_value=0,
                )
                gv: R.Tensor((1, 8, 40, 40, 3), dtype="float32") = lv
                R.output(gv)
            return gv

    graph_model3 = fx.symbolic_trace(InterpolateNHWC3())
    with torch.no_grad():
        mod3 = from_fx(graph_model3, input_info_5d, default_image_layout="NDHWC")
    tvm.ir.assert_structural_equal(mod3, expected_nhwc3)

    class InterpolateNHWC4(Module):
        def forward(self, input):
            return torch.nn.functional.interpolate(
                input,
                size=None,
                scale_factor=(2.0, 4.0, 4.0),
                mode="trilinear",
                align_corners=True,
            )

    @tvm.script.ir_module
    class expected_nhwc4:
        @R.function
        def main(input_5: R.Tensor((1, 4, 10, 10, 3), dtype="float32")) -> R.Tensor(
            (1, 8, 40, 40, 3), dtype="float32"
        ):
            with R.dataflow():
                lv: R.Tensor((1, 8, 40, 40, 3), dtype="float32") = R.image.resize3d(
                    input_5,
                    (8, 40, 40),
                    roi=[0.000000, 0.000000, 0.000000, 0.000000, 0.000000, 0.000000],
                    layout="NDHWC",
                    method="linear",
                    coordinate_transformation_mode="align_corners",
                    rounding_method="",
                    cubic_alpha=-0.75,
                    cubic_exclude=0,
                    extrapolation_value=0,
                )
                gv: R.Tensor((1, 8, 40, 40, 3), dtype="float32") = lv
                R.output(gv)
            return gv

    graph_model4 = fx.symbolic_trace(InterpolateNHWC4())
    with torch.no_grad():
        mod4 = from_fx(graph_model4, input_info_5d, default_image_layout="NDHWC")
    tvm.ir.assert_structural_equal(mod4, expected_nhwc4)


def test_addmm():
    input_info = [
        ([10, 10], "float32"),
        ([10, 10], "float32"),
        ([10, 10], "float32"),
    ]

    class Addmm1(Module):
        def __init__(self):
            super().__init__()

        def forward(self, x1, x2, x3):
            return torch.addmm(x1, x2, x3)

    class Addmm2(Module):
        def __init__(self):
            super().__init__()

        def forward(self, x1, x2, x3):
            return torch.addmm(x1, x2, x3, beta=0.8, alpha=0.5)

    @tvm.script.ir_module
    class expected1:
        @R.function
        def main(
            x1: R.Tensor((10, 10), dtype="float32"),
            x2: R.Tensor((10, 10), dtype="float32"),
            x3: R.Tensor((10, 10), dtype="float32"),
        ) -> R.Tensor((10, 10), dtype="float32"):
            # block 0
            with R.dataflow():
                lv: R.Tensor((10, 10), dtype="float32") = R.matmul(x2, x3, out_dtype="float32")
                lv1: R.Tensor((10, 10), dtype="float32") = R.add(x1, lv)
                gv: R.Tensor((10, 10), dtype="float32") = lv1
                R.output(gv)
            return gv

    @tvm.script.ir_module
    class expected2:
        @R.function
        def main(
            x1: R.Tensor((10, 10), dtype="float32"),
            x2: R.Tensor((10, 10), dtype="float32"),
            x3: R.Tensor((10, 10), dtype="float32"),
        ) -> R.Tensor((10, 10), dtype="float32"):
            # block 0
            with R.dataflow():
                lv: R.Tensor((10, 10), dtype="float32") = R.matmul(x2, x3, out_dtype="float32")
                lv1: R.Tensor((10, 10), dtype="float32") = R.multiply(lv, R.const(0.5, "float32"))
                lv2: R.Tensor((10, 10), dtype="float32") = R.multiply(x1, R.const(0.8, "float32"))
                lv3: R.Tensor((10, 10), dtype="float32") = R.add(lv2, lv1)
                gv: R.Tensor((10, 10), dtype="float32") = lv3
                R.output(gv)
            return gv

    verify_model(Addmm1(), input_info, {}, expected1)
    verify_model(Addmm2(), input_info, {}, expected2)


def test_split():
    input_info = [([1, 3, 10, 10], "float32")]

    class Split1(Module):
        def forward(self, input):
            return torch.split(input, 1, dim=1)

    class Split2(Module):
        def forward(self, input):
            return torch.split(input, [1, 2], dim=1)

    @tvm.script.ir_module
    class expected1:
        @R.function
        def main(input_1: R.Tensor((1, 3, 10, 10), dtype="float32")) -> R.Tuple(
            R.Tensor((1, 1, 10, 10), dtype="float32"),
            R.Tensor((1, 1, 10, 10), dtype="float32"),
            R.Tensor((1, 1, 10, 10), dtype="float32"),
        ):
            # block 0
            with R.dataflow():
                lv: R.Tuple(
                    R.Tensor((1, 1, 10, 10), dtype="float32"),
                    R.Tensor((1, 1, 10, 10), dtype="float32"),
                    R.Tensor((1, 1, 10, 10), dtype="float32"),
                ) = R.split(input_1, indices_or_sections=3, axis=1)
                gv: R.Tuple(
                    R.Tensor((1, 1, 10, 10), dtype="float32"),
                    R.Tensor((1, 1, 10, 10), dtype="float32"),
                    R.Tensor((1, 1, 10, 10), dtype="float32"),
                ) = lv
                R.output(gv)
            return gv

    @tvm.script.ir_module
    class expected2:
        @R.function
        def main(input_1: R.Tensor((1, 3, 10, 10), dtype="float32")) -> R.Tuple(
            R.Tensor((1, 1, 10, 10), dtype="float32"), R.Tensor((1, 2, 10, 10), dtype="float32")
        ):
            # block 0
            with R.dataflow():
                lv: R.Tuple(
                    R.Tensor((1, 1, 10, 10), dtype="float32"),
                    R.Tensor((1, 2, 10, 10), dtype="float32"),
                ) = R.split(input_1, indices_or_sections=[1], axis=1)
                gv: R.Tuple(
                    R.Tensor((1, 1, 10, 10), dtype="float32"),
                    R.Tensor((1, 2, 10, 10), dtype="float32"),
                ) = lv
                R.output(gv)
            return gv

    verify_model(Split1(), input_info, {}, expected1)
    verify_model(Split2(), input_info, {}, expected2)


def test_unbind():
    input_info = [([3, 3, 10, 10], "float32")]

    class Unbind1(Module):
        def forward(self, data):
            return torch.unbind(data)

    class Unbind2(Module):
        def forward(self, data):
            return torch.unbind(data, dim=1)

    @tvm.script.ir_module
    class expected1:
        @R.function
        def main(input_1: R.Tensor((3, 3, 10, 10), dtype="float32")) -> R.Tuple(
            R.Tensor((3, 10, 10), dtype="float32"),
            R.Tensor((3, 10, 10), dtype="float32"),
            R.Tensor((3, 10, 10), dtype="float32"),
        ):
            # block 0
            with R.dataflow():
                lv: R.Tuple(
                    R.Tensor((1, 3, 10, 10), dtype="float32"),
                    R.Tensor((1, 3, 10, 10), dtype="float32"),
                    R.Tensor((1, 3, 10, 10), dtype="float32"),
                ) = R.split(input_1, indices_or_sections=3, axis=0)
                lv1: R.Tensor((1, 3, 10, 10), dtype="float32") = lv[0]
                lv2: R.Tensor((3, 10, 10), dtype="float32") = R.squeeze(lv1, axis=[0])
                lv3: R.Tensor((1, 3, 10, 10), dtype="float32") = lv[1]
                lv4: R.Tensor((3, 10, 10), dtype="float32") = R.squeeze(lv3, axis=[0])
                lv5: R.Tensor((1, 3, 10, 10), dtype="float32") = lv[2]
                lv6: R.Tensor((3, 10, 10), dtype="float32") = R.squeeze(lv5, axis=[0])
                lv7: R.Tuple(
                    R.Tensor((3, 10, 10), dtype="float32"),
                    R.Tensor((3, 10, 10), dtype="float32"),
                    R.Tensor((3, 10, 10), dtype="float32"),
                ) = (lv2, lv4, lv6)
                gv: R.Tuple(
                    R.Tensor((3, 10, 10), dtype="float32"),
                    R.Tensor((3, 10, 10), dtype="float32"),
                    R.Tensor((3, 10, 10), dtype="float32"),
                ) = lv7
                R.output(gv)
            return gv

    @tvm.script.ir_module
    class expected2:
        @R.function
        def main(input_1: R.Tensor((3, 3, 10, 10), dtype="float32")) -> R.Tuple(
            R.Tensor((3, 10, 10), dtype="float32"),
            R.Tensor((3, 10, 10), dtype="float32"),
            R.Tensor((3, 10, 10), dtype="float32"),
        ):
            # block 0
            with R.dataflow():
                lv: R.Tuple(
                    R.Tensor((3, 1, 10, 10), dtype="float32"),
                    R.Tensor((3, 1, 10, 10), dtype="float32"),
                    R.Tensor((3, 1, 10, 10), dtype="float32"),
                ) = R.split(input_1, indices_or_sections=3, axis=1)
                lv1: R.Tensor((3, 1, 10, 10), dtype="float32") = lv[0]
                lv2: R.Tensor((3, 10, 10), dtype="float32") = R.squeeze(lv1, axis=[1])
                lv3: R.Tensor((3, 1, 10, 10), dtype="float32") = lv[1]
                lv4: R.Tensor((3, 10, 10), dtype="float32") = R.squeeze(lv3, axis=[1])
                lv5: R.Tensor((3, 1, 10, 10), dtype="float32") = lv[2]
                lv6: R.Tensor((3, 10, 10), dtype="float32") = R.squeeze(lv5, axis=[1])
                lv7: R.Tuple(
                    R.Tensor((3, 10, 10), dtype="float32"),
                    R.Tensor((3, 10, 10), dtype="float32"),
                    R.Tensor((3, 10, 10), dtype="float32"),
                ) = (lv2, lv4, lv6)
                gv: R.Tuple(
                    R.Tensor((3, 10, 10), dtype="float32"),
                    R.Tensor((3, 10, 10), dtype="float32"),
                    R.Tensor((3, 10, 10), dtype="float32"),
                ) = lv7
                R.output(gv)
            return gv

    verify_model(Unbind1(), input_info, {}, expected1)
    verify_model(Unbind2(), input_info, {}, expected2)


def test_cumsum():
    input_info = [([1, 2, 3, 4], "float32")]

    class Cumsum(Module):
        def forward(self, input):
            return torch.cumsum(input, dim=1, dtype=torch.int32)

    @tvm.script.ir_module
    class expected1:
        @R.function
        def main(input_1: R.Tensor((1, 2, 3, 4), dtype="float32")) -> R.Tensor(
            (1, 2, 3, 4), dtype="int32"
        ):
            # block 0
            with R.dataflow():
                lv: R.Tensor((1, 2, 3, 4), dtype="int32") = R.cumsum(input_1, axis=1, dtype="int32")
                gv: R.Tensor((1, 2, 3, 4), dtype="int32") = lv
                R.output(gv)
            return gv

    verify_model(Cumsum(), input_info, {}, expected1)


def test_chunk():
    input_info = [([1, 3, 10, 10], "float32")]

    class Chunk(Module):
        def forward(self, input):
            return torch.chunk(input, 3, dim=1)

    @tvm.script.ir_module
    class Expected:
        @R.function
        def main(input_1: R.Tensor((1, 3, 10, 10), dtype="float32")) -> R.Tuple(
            R.Tensor((1, 1, 10, 10), dtype="float32"),
            R.Tensor((1, 1, 10, 10), dtype="float32"),
            R.Tensor((1, 1, 10, 10), dtype="float32"),
        ):
            # block 0
            with R.dataflow():
                lv: R.Tuple(
                    R.Tensor((1, 1, 10, 10), dtype="float32"),
                    R.Tensor((1, 1, 10, 10), dtype="float32"),
                    R.Tensor((1, 1, 10, 10), dtype="float32"),
                ) = R.split(input_1, indices_or_sections=3, axis=1)
                gv: R.Tuple(
                    R.Tensor((1, 1, 10, 10), dtype="float32"),
                    R.Tensor((1, 1, 10, 10), dtype="float32"),
                    R.Tensor((1, 1, 10, 10), dtype="float32"),
                ) = lv
                R.output(gv)
            return gv

    verify_model(Chunk(), input_info, {}, Expected)


def test_inplace_fill():
    class InplaceFill(Module):
        def forward(self, input):
            input.fill_(1.5)
            return input

    @tvm.script.ir_module
    class Expected:
        @R.function
        def main(inp_0: R.Tensor((10, 10), dtype="float32")) -> R.Tensor((10, 10), dtype="float32"):
            with R.dataflow():
                lv: R.Tensor((10, 10), dtype="float32") = R.full(
                    R.shape([10, 10]), R.const(1.5, "float32"), dtype="float32"
                )
                gv: R.Tensor((10, 10), dtype="float32") = lv
                R.output(gv)
            return gv

    verify_model(InplaceFill(), [([10, 10], "float32")], {}, Expected)


def test_masked_fill_inplace():
    class Masked_Fill_Inplace(Module):
        def forward(self, input: torch.Tensor, mask: torch.Tensor):
            input.masked_fill_(mask, 1.5)
            return input

    @tvm.script.ir_module
    class Expected:
        @R.function
        def main(
            input: R.Tensor((10, 10), dtype="float32"), mask: R.Tensor((10, 10), dtype="bool")
        ) -> R.Tensor((10, 10), dtype="float32"):
            with R.dataflow():
                lv: R.Tensor((10, 10), dtype="float32") = R.full_like(
                    input, R.const(1.5, "float32")
                )
                lv1: R.Tensor((10, 10), dtype="float32") = R.where(mask, lv, input)
                gv: R.Tensor((10, 10), dtype="float32") = lv1
                R.output(gv)
            return gv

    input_info = [((10, 10), "float32"), ((10, 10), "bool")]
    verify_model(Masked_Fill_Inplace(), input_info, {}, Expected)


def test_arange():
    import numpy as np

    torch.set_grad_enabled(False)
    torch.random.manual_seed(0)

    class Arange(Module):
        def forward(self, input):
            return torch.arange(0, 20, dtype=torch.int32)

    graph_model = fx.symbolic_trace(Arange())
    mod = from_fx(graph_model, [([10, 10], "float32")])
    assert len(mod["main"].body.blocks) == 1
    assert len(mod["main"].body.blocks[0].bindings) == 1
    assert isinstance(mod["main"].body.blocks[0].bindings[0].value, relax.Constant)
    tvm.testing.assert_allclose(
        mod["main"].body.blocks[0].bindings[0].value.data.numpy(),
        np.arange(0, 20, dtype="int32"),
    )


def test_empty():
    class Empty(Module):
        def forward(self, input):
            return torch.empty((10, 10), dtype=torch.float32)

    graph_model = fx.symbolic_trace(Empty())
    mod = from_fx(graph_model, [([10, 10], "float32")])
    assert len(mod["main"].body.blocks) == 1
    assert len(mod["main"].body.blocks[0].bindings) == 1
    assert isinstance(mod["main"].body.blocks[0].bindings[0].value, relax.Constant)
    assert mod["main"].body.blocks[0].bindings[0].value.data.shape == (10, 10)
    assert mod["main"].body.blocks[0].bindings[0].value.data.dtype == "float32"


def test_tensor():
    class Empty1(Module):
        def forward(self, input):
            return torch.tensor(3, dtype=torch.float32)

    class Empty2(Module):
        def forward(self, input):
            return torch.tensor(3)

    graph_model1 = fx.symbolic_trace(Empty1())
    mod1 = from_fx(graph_model1, [([10, 10], "float32")])
    assert len(mod1["main"].body.blocks) == 1
    assert len(mod1["main"].body.blocks[0].bindings) == 1
    assert isinstance(mod1["main"].body.blocks[0].bindings[0].value, relax.Constant)
    assert mod1["main"].body.blocks[0].bindings[0].value.data.shape == ()
    assert mod1["main"].body.blocks[0].bindings[0].value.data.dtype == "float32"

    graph_model2 = fx.symbolic_trace(Empty2())
    mod2 = from_fx(graph_model2, [([10, 10], "float32")])
    assert len(mod2["main"].body.blocks) == 1
    assert len(mod2["main"].body.blocks[0].bindings) == 1
    assert isinstance(mod2["main"].body.blocks[0].bindings[0].value, relax.Constant)
    assert mod2["main"].body.blocks[0].bindings[0].value.data.shape == ()
    assert mod2["main"].body.blocks[0].bindings[0].value.data.dtype == "int64"


def test_new_ones():
    input_info = [([1, 2, 3], "float32")]

    class NewOnes(Module):
        def forward(self, x):
            return x.new_ones(1, 2, 3)

    @tvm.script.ir_module
    class expected1:
        @R.function
        def main(x: R.Tensor((1, 2, 3), dtype="float32")) -> R.Tensor((1, 2, 3), dtype="float32"):
            # block 0
            with R.dataflow():
                lv: R.Tensor((1, 2, 3), dtype="float32") = R.full(
                    (1, 2, 3), R.const(1, "float32"), dtype="float32"
                )
                gv: R.Tensor((1, 2, 3), dtype="float32") = lv
                R.output(gv)
            return gv

    verify_model(NewOnes(), input_info, {}, expected1)


def test_new_zeros():
    input_info = [([1, 128, 128], "float32")]

    class NewZeros(Module):
        def forward(self, x):
            return x.new_zeros(1, 128, 128)

    @tvm.script.ir_module
    class expected:
        @R.function
        def main(x: R.Tensor((1, 128, 128), dtype="float32")) -> R.Tensor(
            (1, 128, 128), dtype="float32"
        ):
            # block 0
            with R.dataflow():
                lv: R.Tensor((1, 128, 128), dtype="float32") = R.full(
                    (1, 128, 128), R.const(0.0, "float32"), dtype="float32"
                )
                gv: R.Tensor((1, 128, 128), dtype="float32") = lv
                R.output(gv)
            return gv

    verify_model(NewZeros(), input_info, {}, expected)


def test_expand():
    input_info = [([1, 2, 3, 4], "float32")]

    class Expand1(Module):
        def forward(self, x):
            return x.expand(4, 2, 3, 4)

    class Expand2(Module):
        def forward(self, x):
            return x.expand(4, -1, -1, 4)

    @tvm.script.ir_module
    class expected1:
        @R.function
        def main(x: R.Tensor((1, 2, 3, 4), dtype="float32")) -> R.Tensor(
            (4, 2, 3, 4), dtype="float32"
        ):
            # block 0
            with R.dataflow():
                lv: R.Tensor((4, 2, 3, 4), dtype="float32") = R.broadcast_to(x, (4, 2, 3, 4))
                gv: R.Tensor((4, 2, 3, 4), dtype="float32") = lv
                R.output(gv)
            return gv

    verify_model(Expand1(), input_info, {}, expected1)
    verify_model(Expand2(), input_info, {}, expected1)


def test_reduce():
    input_info = [([1, 2, 3, 4], "float32")]

    # sum
    class Sum(Module):
        def forward(self, x):
            return torch.sum(x, (2, 1))

    @tvm.script.ir_module
    class expected1:
        @R.function
        def main(inp_0: R.Tensor((1, 2, 3, 4), dtype="float32")) -> R.Tensor(
            (1, 4), dtype="float32"
        ):
            # block 0
            with R.dataflow():
                lv: R.Tensor((1, 4), dtype="float32") = R.sum(inp_0, axis=[2, 1], keepdims=False)
                gv: R.Tensor((1, 4), dtype="float32") = lv
                R.output(gv)
            return gv

    verify_model(Sum(), input_info, {}, expected1)


def test_datatype():
    input_info = [([1, 2, 3, 4], "float32")]

    # float
    class ToFloat(Module):
        def forward(self, x):
            return x.float()

    @tvm.script.ir_module
    class expected1:
        @R.function
        def main(x: R.Tensor((1, 2, 3, 4), dtype="float32")) -> R.Tensor(
            (1, 2, 3, 4), dtype="float32"
        ):
            # block 0
            with R.dataflow():
                lv: R.Tensor((1, 2, 3, 4), dtype="float32") = R.astype(x, dtype="float32")
                gv: R.Tensor((1, 2, 3, 4), dtype="float32") = lv
                R.output(gv)
            return gv

    verify_model(ToFloat(), input_info, {}, expected1)

    # half
    class ToHalf(Module):
        def forward(self, x):
            return x.half()

    @tvm.script.ir_module
    class expected2:
        @R.function
        def main(x: R.Tensor((1, 2, 3, 4), dtype="float32")) -> R.Tensor(
            (1, 2, 3, 4), dtype="float16"
        ):
            # block 0
            with R.dataflow():
                lv: R.Tensor((1, 2, 3, 4), dtype="float16") = R.astype(x, dtype="float16")
                gv: R.Tensor((1, 2, 3, 4), dtype="float16") = lv
                R.output(gv)
            return gv

    verify_model(ToHalf(), input_info, {}, expected2)

    # type
    class Type(Module):
        def forward(self, x):
            return x.type(torch.float32)

    # type
    class TypeFromAttr(Module):
        def forward(self, x):
            return x.type(x.getattr("dtype"))

    # astype
    class AsType(Module):
        def forward(self, x):
            return x.astype(torch.float32)

    verify_model(Type(), input_info, {}, expected1)
    verify_model(TypeFromAttr(), input_info, {}, expected1)
    verify_model(AsType(), input_info, {}, expected1)


def test_meshgrid():
    input_infos = [
        (
            [
                3,
            ],
            "float32",
        ),
        (
            [
                3,
            ],
            "float32",
        ),
    ]

    class Meshgrid1(Module):
        def forward(self, input1, input2):
            return torch.meshgrid((input1, input2), indexing="ij")

    class Meshgrid2(Module):
        def forward(self, input1, input2):
            return torch.meshgrid((input1, input2), indexing="xy")

    @tvm.script.ir_module
    class expected1:
        @R.function
        def main(
            inp_0: R.Tensor((3,), dtype="float32"), inp_1: R.Tensor((3,), dtype="float32")
        ) -> R.Tuple(R.Tensor((3, 3), dtype="float32"), R.Tensor((3, 3), dtype="float32")):
            with R.dataflow():
                lv: R.Tuple(
                    R.Tensor((3, 3), dtype="float32"), R.Tensor((3, 3), dtype="float32")
                ) = R.meshgrid((inp_0, inp_1), indexing="ij")
                gv: R.Tuple(
                    R.Tensor((3, 3), dtype="float32"), R.Tensor((3, 3), dtype="float32")
                ) = lv
                R.output(gv)
            return gv

    @tvm.script.ir_module
    class expected2:
        @R.function
        def main(
            inp_0: R.Tensor((3,), dtype="float32"), inp_1: R.Tensor((3,), dtype="float32")
        ) -> R.Tuple(R.Tensor((3, 3), dtype="float32"), R.Tensor((3, 3), dtype="float32")):
            with R.dataflow():
                lv: R.Tuple(
                    R.Tensor((3, 3), dtype="float32"), R.Tensor((3, 3), dtype="float32")
                ) = R.meshgrid((inp_0, inp_1), indexing="xy")
                gv: R.Tuple(
                    R.Tensor((3, 3), dtype="float32"), R.Tensor((3, 3), dtype="float32")
                ) = lv
                R.output(gv)
            return gv

    verify_model(Meshgrid1(), input_infos, {}, expected1)
    verify_model(Meshgrid2(), input_infos, {}, expected2)


def test_permute():
    input_info = [([1, 2, 3, 4], "float32")]

    class Permute1(Module):
        def forward(self, x):
            return x.permute(0, 3, 2, 1)

    class Permute2(Module):
        def forward(self, x):
            return torch.permute(x, (0, 3, 2, 1))

    @tvm.script.ir_module
    class expected1:
        @R.function
        def main(x: R.Tensor((1, 2, 3, 4), dtype="float32")) -> R.Tensor(
            (1, 4, 3, 2), dtype="float32"
        ):
            # block 0
            with R.dataflow():
                lv: R.Tensor((1, 4, 3, 2), dtype="float32") = R.permute_dims(x, axes=[0, 3, 2, 1])
                gv: R.Tensor((1, 4, 3, 2), dtype="float32") = lv
                R.output(gv)
            return gv

    verify_model(Permute1(), input_info, {}, expected1)
    verify_model(Permute2(), input_info, {}, expected1)


def test_reshape():
    input_info = [([1, 2, 3, 4], "float32")]

    class Reshape(Module):
        def forward(self, x):
            return x.reshape(2, 12)

    @tvm.script.ir_module
    class expected1:
        @R.function
        def main(x: R.Tensor((1, 2, 3, 4), dtype="float32")) -> R.Tensor((2, 12), dtype="float32"):
            # block 0
            with R.dataflow():
                lv: R.Tensor((2, 12), dtype="float32") = R.reshape(x, (2, 12))
                gv: R.Tensor((2, 12), dtype="float32") = lv
                R.output(gv)
            return gv

    verify_model(Reshape(), input_info, {}, expected1)


def test_tile():
    input_info = [([1, 3], "float32")]

    class Tile1(Module):
        def forward(self, x):
            return x.tile((2,))

    class Tile2(Module):
        def forward(self, x):
            return x.tile(4, 2)

    class Tile3(Module):
        def forward(self, x):
            return torch.tile(x, (4, 2))

    @tvm.script.ir_module
    class expected1:
        @R.function
        def main(x: R.Tensor((1, 3), dtype="float32")) -> R.Tensor((1, 6), dtype="float32"):
            # block 0
            with R.dataflow():
                lv: R.Tensor((1, 6), dtype="float32") = R.tile(x, [2])
                gv: R.Tensor((1, 6), dtype="float32") = lv
                R.output(gv)
            return gv

    @tvm.script.ir_module
    class expected2:
        @R.function
        def main(x: R.Tensor((1, 3), dtype="float32")) -> R.Tensor((4, 6), dtype="float32"):
            # block 0
            with R.dataflow():
                lv: R.Tensor((4, 6), dtype="float32") = R.tile(x, [4, 2])
                gv: R.Tensor((4, 6), dtype="float32") = lv
                R.output(gv)
            return gv

    verify_model(Tile1(), input_info, {}, expected1)
    verify_model(Tile2(), input_info, {}, expected2)
    verify_model(Tile3(), input_info, {}, expected2)


def test_transpose():
    input_info = [([1, 2, 3, 4], "float32")]

    class Transpose(Module):
        def forward(self, x):
            return x.transpose(1, 3)

    @tvm.script.ir_module
    class expected1:
        @R.function
        def main(x: R.Tensor((1, 2, 3, 4), dtype="float32")) -> R.Tensor(
            (1, 4, 3, 2), dtype="float32"
        ):
            # block 0
            with R.dataflow():
                lv: R.Tensor((1, 4, 3, 2), dtype="float32") = R.permute_dims(x, axes=[0, 3, 2, 1])
                gv: R.Tensor((1, 4, 3, 2), dtype="float32") = lv
                R.output(gv)
            return gv

    verify_model(Transpose(), input_info, {}, expected1)


def test_repeat():
    class Tile1(Module):
        def forward(self, x: torch.Tensor):
            return x.repeat(2)

    class Tile2(Module):
        def forward(self, x: torch.Tensor):
            return x.repeat(4, 2)

    @tvm.script.ir_module
    class expected1:
        @R.function
        def main(x: R.Tensor((3,), dtype="float32")) -> R.Tensor((6,), dtype="float32"):
            # block 0
            with R.dataflow():
                lv: R.Tensor((6,), dtype="float32") = R.tile(x, 2)
                gv: R.Tensor((6,), dtype="float32") = lv
                R.output(gv)
            return gv

    @tvm.script.ir_module
    class expected2:
        @R.function
        def main(x: R.Tensor((1, 3), dtype="float32")) -> R.Tensor((4, 6), dtype="float32"):
            # block 0
            with R.dataflow():
                lv: R.Tensor((4, 6), dtype="float32") = R.tile(x, [4, 2])
                gv: R.Tensor((4, 6), dtype="float32") = lv
                R.output(gv)
            return gv

    verify_model(Tile1(), [([3], "float32")], {}, expected1)
    verify_model(Tile2(), [([1, 3], "float32")], {}, expected2)
    verify_model(Tile2(), [(torch.Size([1, 3]), "float32")], {}, expected2)


def test_roll():
    class Roll1(Module):
        def forward(self, x):
            return torch.roll(x, 1)

    class Roll2(Module):
        def forward(self, x):
            return torch.roll(x, -1, 0)

    class Roll3(Module):
        def forward(self, x):
            return torch.roll(x, shifts=(2, 1), dims=(0, 1))

    # Test case 1: torch.roll(x, 1)
    @I.ir_module
    class Expected1:
        @R.function
        def main(inp_0: R.Tensor((4, 2), dtype="int64")) -> R.Tensor((4, 2), dtype="int64"):
            with R.dataflow():
                lv: R.Tensor((8,), dtype="int64") = R.reshape(inp_0, R.shape([8]))
                lv1: R.Tensor((7,), dtype="int64") = R.strided_slice(
                    lv,
                    axes=[0],
                    begin=[R.prim_value(0)],
                    end=[R.prim_value(7)],
                    strides=[R.prim_value(1)],
                    assume_inbound=False,
                )
                lv2: R.Tensor((1,), dtype="int64") = R.strided_slice(
                    lv,
                    axes=[0],
                    begin=[R.prim_value(7)],
                    end=[R.prim_value(8)],
                    strides=[R.prim_value(1)],
                    assume_inbound=False,
                )
                lv3: R.Tensor((8,), dtype="int64") = R.concat((lv2, lv1), axis=0)
                lv4: R.Tensor((4, 2), dtype="int64") = R.reshape(lv3, R.shape([4, 2]))
                gv: R.Tensor((4, 2), dtype="int64") = lv4
                R.output(gv)
            return gv

    # Test case 2: torch.roll(x, -1, 0)
    @I.ir_module
    class Expected2:
        @R.function
        def main(inp_0: R.Tensor((4, 2), dtype="int64")) -> R.Tensor((4, 2), dtype="int64"):
            with R.dataflow():
                lv: R.Tensor((1, 2), dtype="int64") = R.strided_slice(
                    inp_0,
                    axes=[0],
                    begin=[R.prim_value(0)],
                    end=[R.prim_value(1)],
                    strides=[R.prim_value(1)],
                    assume_inbound=False,
                )
                lv1: R.Tensor((3, 2), dtype="int64") = R.strided_slice(
                    inp_0,
                    axes=[0],
                    begin=[R.prim_value(1)],
                    end=[R.prim_value(4)],
                    strides=[R.prim_value(1)],
                    assume_inbound=False,
                )
                lv2: R.Tensor((4, 2), dtype="int64") = R.concat((lv1, lv), axis=0)
                gv: R.Tensor((4, 2), dtype="int64") = lv2
                R.output(gv)
            return gv

    # Test case 3: torch.roll(x, shifts=(2, 1), dims=(0, 1))
    @I.ir_module
    class Expected3:
        @R.function
        def main(inp_0: R.Tensor((4, 2), dtype="int64")) -> R.Tensor((4, 2), dtype="int64"):
            with R.dataflow():
                lv: R.Tensor((2, 2), dtype="int64") = R.strided_slice(
                    inp_0,
                    axes=[0],
                    begin=[R.prim_value(0)],
                    end=[R.prim_value(2)],
                    strides=[R.prim_value(1)],
                    assume_inbound=False,
                )
                lv1: R.Tensor((2, 2), dtype="int64") = R.strided_slice(
                    inp_0,
                    axes=[0],
                    begin=[R.prim_value(2)],
                    end=[R.prim_value(4)],
                    strides=[R.prim_value(1)],
                    assume_inbound=False,
                )
                lv2: R.Tensor((4, 2), dtype="int64") = R.concat((lv1, lv), axis=0)
                lv3: R.Tensor((4, 1), dtype="int64") = R.strided_slice(
                    lv2,
                    axes=[1],
                    begin=[R.prim_value(0)],
                    end=[R.prim_value(1)],
                    strides=[R.prim_value(1)],
                    assume_inbound=False,
                )
                lv4: R.Tensor((4, 1), dtype="int64") = R.strided_slice(
                    lv2,
                    axes=[1],
                    begin=[R.prim_value(1)],
                    end=[R.prim_value(2)],
                    strides=[R.prim_value(1)],
                    assume_inbound=False,
                )
                lv5: R.Tensor((4, 2), dtype="int64") = R.concat((lv4, lv3), axis=1)
                gv: R.Tensor((4, 2), dtype="int64") = lv5
                R.output(gv)
            return gv

    input_info = [([4, 2], "int64")]

    verify_model(Roll1(), input_info, {}, Expected1)
    verify_model(Roll2(), input_info, {}, Expected2)
    verify_model(Roll3(), input_info, {}, Expected3)


def test_view():
    input_info = [([1, 2, 3, 4], "float32")]

    class View(Module):
        def forward(self, x):
            return x.view(2, 12)

    @tvm.script.ir_module
    class expected1:
        @R.function
        def main(x: R.Tensor((1, 2, 3, 4), dtype="float32")) -> R.Tensor((2, 12), dtype="float32"):
            # block 0
            with R.dataflow():
                lv: R.Tensor((2, 12), dtype="float32") = R.reshape(x, (2, 12))
                gv: R.Tensor((2, 12), dtype="float32") = lv
                R.output(gv)
            return gv

    verify_model(View(), input_info, {}, expected1)


def test_keep_params():
    class Conv2D1(Module):
        def __init__(self):
            super().__init__()
            self.conv = torch.nn.Conv2d(3, 6, 7, bias=True)

        def forward(self, input):
            return self.conv(input)

    @tvm.script.ir_module
    class expected1:
        @R.function
        def main(
            input_1: R.Tensor((1, 3, 10, 10), dtype="float32"),
            w1: R.Tensor((6,), dtype="float32"),
            w2: R.Tensor((6, 3, 7, 7), dtype="float32"),
        ) -> R.Tensor((1, 6, 4, 4), dtype="float32"):
            R.func_attr({"num_input": 1})
            # block 0
            with R.dataflow():
                lv1: R.Tensor((1, 6, 4, 4), dtype="float32") = R.nn.conv2d(
                    input_1,
                    w2,
                    strides=[1, 1],
                    padding=[0, 0, 0, 0],
                    dilation=[1, 1],
                    data_layout="NCHW",
                    kernel_layout="OIHW",
                    out_layout="NCHW",
                    out_dtype="float32",
                )
                lv2: R.Tensor((1, 6, 1, 1), dtype="float32") = R.reshape(w1, [1, 6, 1, 1])
                lv3: R.Tensor((1, 6, 4, 4), dtype="float32") = R.add(lv1, lv2)
                gv: R.Tensor((1, 6, 4, 4), dtype="float32") = lv3
                R.output(gv)
            return gv

    model = Conv2D1()
    graph_model = fx.symbolic_trace(model)
    mod = from_fx(graph_model, [([1, 3, 10, 10], "float32")], keep_params_as_input=True)
    mod, params = detach_params(mod)
    tvm.ir.assert_structural_equal(mod, expected1)
    func = mod["main"]
    params = params["main"]

    assert len(params) == len(func.params) - 1
    for param_var, param_tensor in zip(func.params[1:], params):
        assert tuple(x.value for x in param_var.ty.shape.values) == param_tensor.shape
        assert param_var.ty.dtype == param_tensor.dtype

    tvm.testing.assert_allclose(params[0].numpy(), model.conv.bias.detach().detach().numpy())
    tvm.testing.assert_allclose(params[1].numpy(), model.conv.weight.detach().detach().numpy())


def test_unwrap_unit_return_tuple():
    class Identity(Module):
        def __init__(self):
            super().__init__()

        def forward(self, x):
            return (x,)

    @tvm.script.ir_module
    class Expected:
        @R.function
        def main(inp_0: R.Tensor((256, 256), dtype="float32")) -> R.Tensor(
            (256, 256), dtype="float32"
        ):
            with R.dataflow():
                gv: R.Tensor((256, 256), dtype="float32") = inp_0
                R.output(gv)
            return gv

    graph_model = fx.symbolic_trace(Identity())
    mod = from_fx(graph_model, [([256, 256], "float32")], unwrap_unit_return_tuple=True)
    tvm.ir.assert_structural_equal(mod, Expected)


def test_no_bind_return_tuple():
    class Identity(Module):
        def __init__(self):
            super().__init__()

        def forward(self, x, y):
            return (x, y)

    @tvm.script.ir_module
    class Expected:
        @R.function
        def main(
            inp_0: R.Tensor((256, 256), dtype="float32"),
            inp_1: R.Tensor((256, 256), dtype="float32"),
        ) -> R.Tuple(R.Tensor((256, 256), dtype="float32"), R.Tensor((256, 256), dtype="float32")):
            with R.dataflow():
                gv: R.Tensor((256, 256), dtype="float32") = inp_0
                gv1: R.Tensor((256, 256), dtype="float32") = inp_1
                R.output(gv, gv1)
            return (gv, gv1)

    graph_model = fx.symbolic_trace(Identity())
    mod = from_fx(
        graph_model, [([256, 256], "float32"), ([256, 256], "float32")], no_bind_return_tuple=True
    )
    tvm.ir.assert_structural_equal(mod, Expected)


def test_argmax():
    class Argmax1(Module):
        def __init__(self) -> None:
            super().__init__()

        def forward(self, input):
            return torch.argmax(input, dim=-1)

    class Argmax2(Module):
        def __init__(self) -> None:
            super().__init__()

        def forward(self, input):
            return torch.argmax(input, dim=-1, keepdim=True)

    @tvm.script.ir_module
    class Expected1:
        @R.function
        def main(inp_0: R.Tensor((256, 256), dtype="float32")) -> R.Tensor((256,), dtype="int64"):
            with R.dataflow():
                lv: R.Tensor((256,), dtype="int64") = R.argmax(inp_0, axis=-1, keepdims=False)
                gv: R.Tensor((256,), dtype="int64") = lv
                R.output(gv)
            return gv

    @tvm.script.ir_module
    class Expected2:
        @R.function
        def main(inp_0: R.Tensor((256, 256), dtype="float32")) -> R.Tensor((256, 1), dtype="int64"):
            with R.dataflow():
                lv: R.Tensor((256, 1), dtype="int64") = R.argmax(inp_0, axis=-1, keepdims=True)
                gv: R.Tensor((256, 1), dtype="int64") = lv
                R.output(gv)
            return gv

    verify_model(Argmax1(), [([256, 256], "float32")], {}, Expected1)
    verify_model(Argmax2(), [([256, 256], "float32")], {}, Expected2)


def test_argmin():
    class Argmin1(Module):
        def __init__(self) -> None:
            super().__init__()

        def forward(self, input):
            return torch.argmin(input)

    class Argmin2(Module):
        def __init__(self) -> None:
            super().__init__()

        def forward(self, input):
            return torch.argmin(input, keepdim=True)

    @tvm.script.ir_module
    class Expected1:
        @R.function
        def main(inp_0: R.Tensor((256, 256), dtype="float32")) -> R.Tensor((), dtype="int64"):
            with R.dataflow():
                lv: R.Tensor((), dtype="int64") = R.argmin(inp_0, axis=None, keepdims=False)
                gv: R.Tensor((), dtype="int64") = lv
                R.output(gv)
            return gv

    @tvm.script.ir_module
    class Expected2:
        @R.function
        def main(inp_0: R.Tensor((256, 256), dtype="float32")) -> R.Tensor((1, 1), dtype="int64"):
            with R.dataflow():
                lv: R.Tensor((1, 1), dtype="int64") = R.argmin(inp_0, axis=None, keepdims=True)
                gv: R.Tensor((1, 1), dtype="int64") = lv
                R.output(gv)
            return gv

    verify_model(Argmin1(), [([256, 256], "float32")], {}, Expected1)
    verify_model(Argmin2(), [([256, 256], "float32")], {}, Expected2)


def test_to():
    class To1(Module):
        def forward(self, input):
            return input.to(torch.float16)

    class To2(Module):
        def forward(self, input):
            return input.to("cpu")

    @I.ir_module
    class Expected1:
        @R.function
        def main(inp_0: R.Tensor((256, 256), dtype="float32")) -> R.Tensor(
            (256, 256), dtype="float16"
        ):
            with R.dataflow():
                lv: R.Tensor((256, 256), dtype="float16") = R.astype(inp_0, dtype="float16")
                gv: R.Tensor((256, 256), dtype="float16") = lv
                R.output(gv)
            return gv

    @I.ir_module
    class Expected2:
        @R.function
        def main(inp_0: R.Tensor((256, 256), dtype="float32")) -> R.Tensor(
            (256, 256), dtype="float32"
        ):
            with R.dataflow():
                gv: R.Tensor((256, 256), dtype="float32") = inp_0
                R.output(gv)
            return gv

    verify_model(To1(), [([256, 256], "float32")], {}, Expected1)
    verify_model(To2(), [([256, 256], "float32")], {}, Expected2)


def test_mean():
    class Mean(Module):
        def forward(self, input):
            return input.mean(-1)

    class MeanKeepDim(Module):
        def forward(self, input):
            return input.mean(-1, keepdim=True)

    @I.ir_module
    class Expected1:
        @R.function
        def main(inp_0: R.Tensor((256, 256), dtype="float32")) -> R.Tensor((256,), dtype="float32"):
            with R.dataflow():
                lv: R.Tensor((256,), dtype="float32") = R.mean(inp_0, axis=[-1], keepdims=False)
                gv: R.Tensor((256,), dtype="float32") = lv
                R.output(gv)
            return gv

    @I.ir_module
    class Expected2:
        @R.function
        def main(inp_0: R.Tensor((256, 256), dtype="float32")) -> R.Tensor(
            (256, 1), dtype="float32"
        ):
            with R.dataflow():
                lv: R.Tensor((256, 1), dtype="float32") = R.mean(inp_0, axis=[-1], keepdims=True)
                gv: R.Tensor((256, 1), dtype="float32") = lv
                R.output(gv)
            return gv

    verify_model(Mean(), [([256, 256], "float32")], {}, Expected1)
    verify_model(MeanKeepDim(), [([256, 256], "float32")], {}, Expected2)


def test_cat():
    class Cat0(Module):
        def forward(self, x, y):
            return torch.cat((x, y))

    class Cat1(Module):
        def forward(self, x, y):
            return torch.cat((x, y), dim=1)

    class Cat2(Module):
        def forward(self, x, y):
            return torch.cat((x, y), 1)

    class Cat3(Module):
        def forward(self, x, y):
            return torch.concat((x, y), dim=0)

    @I.ir_module
    class Expected1:
        @R.function
        def main(
            inp_0: R.Tensor((2, 3), dtype="float32"),
            inp_1: R.Tensor((2, 3), dtype="float32"),
        ) -> R.Tensor((4, 3), dtype="float32"):
            with R.dataflow():
                lv: R.Tensor((4, 3), dtype="float32") = R.concat((inp_0, inp_1), axis=0)
                gv: R.Tensor((4, 3), dtype="float32") = lv
                R.output(gv)
            return gv

    @I.ir_module
    class Expected2:
        @R.function
        def main(
            inp_0: R.Tensor((2, 3), dtype="float32"),
            inp_1: R.Tensor((2, 3), dtype="float32"),
        ) -> R.Tensor((2, 6), dtype="float32"):
            with R.dataflow():
                lv: R.Tensor((2, 6), dtype="float32") = R.concat((inp_0, inp_1), axis=1)
                gv: R.Tensor((2, 6), dtype="float32") = lv
                R.output(gv)
            return gv

    verify_model(Cat0(), [([2, 3], "float32"), ([2, 3], "float32")], {}, Expected1)
    verify_model(Cat1(), [([2, 3], "float32"), ([2, 3], "float32")], {}, Expected2)
    verify_model(Cat2(), [([2, 3], "float32"), ([2, 3], "float32")], {}, Expected2)
    verify_model(Cat3(), [([2, 3], "float32"), ([2, 3], "float32")], {}, Expected1)


def test_max():
    class Max(Module):
        def forward(self, x, y):
            return torch.max(x, y)

    @I.ir_module
    class Expected1:
        @R.function
        def main(
            inp_0: R.Tensor((256, 256), dtype="float32"),
            inp_1: R.Tensor((256, 256), dtype="float32"),
        ) -> R.Tensor((256, 256), dtype="float32"):
            with R.dataflow():
                lv: R.Tensor((256, 256), dtype="float32") = R.maximum(inp_0, inp_1)
                gv: R.Tensor((256, 256), dtype="float32") = lv
                R.output(gv)
            return gv

    verify_model(Max(), [([256, 256], "float32"), ([256, 256], "float32")], {}, Expected1)


def test_min():
    class Min(Module):
        def forward(self, x, y):
            return torch.min(x, y)

    @I.ir_module
    class Expected1:
        @R.function
        def main(
            inp_0: R.Tensor((256, 256), dtype="float32"),
            inp_1: R.Tensor((256, 256), dtype="float32"),
        ) -> R.Tensor((256, 256), dtype="float32"):
            with R.dataflow():
                lv: R.Tensor((256, 256), dtype="float32") = R.minimum(inp_0, inp_1)
                gv: R.Tensor((256, 256), dtype="float32") = lv
                R.output(gv)
            return gv

    verify_model(Min(), [([256, 256], "float32"), ([256, 256], "float32")], {}, Expected1)


def test_atan2():
    class Atan2(Module):
        def forward(self, x, y):
            return torch.atan2(x, y)

    @I.ir_module
    class Expected1:
        @R.function
        def main(
            inp_0: R.Tensor((256, 256), dtype="float32"),
            inp_1: R.Tensor((256, 256), dtype="float32"),
        ) -> R.Tensor((256, 256), dtype="float32"):
            with R.dataflow():
                lv: R.Tensor((256, 256), dtype="float32") = R.atan2(inp_0, inp_1)
                gv: R.Tensor((256, 256), dtype="float32") = lv
                R.output(gv)
            return gv

    verify_model(Atan2(), [([256, 256], "float32"), ([256, 256], "float32")], {}, Expected1)


def test_attention():
    @I.ir_module
    class Expected1:
        @R.function
        def main(
            inp_0: R.Tensor((32, 8, 128, 64), dtype="float32"),
            inp_1: R.Tensor((32, 8, 128, 64), dtype="float32"),
            inp_2: R.Tensor((32, 8, 128, 64), dtype="float32"),
        ) -> R.Tensor((32, 8, 128, 64), dtype="float32"):
            with R.dataflow():
                lv: R.Tensor((32, 128, 8, 64), dtype="float32") = R.permute_dims(
                    inp_0, axes=[0, 2, 1, 3]
                )
                lv1: R.Tensor((32, 128, 8, 64), dtype="float32") = R.permute_dims(
                    inp_1, axes=[0, 2, 1, 3]
                )
                lv2: R.Tensor((32, 128, 8, 64), dtype="float32") = R.permute_dims(
                    inp_2, axes=[0, 2, 1, 3]
                )
                lv3: R.Tensor((32, 128, 8, 64), dtype="float32") = R.nn.attention(
                    lv, lv1, lv2, scale=None
                )
                lv4: R.Tensor((32, 8, 128, 64), dtype="float32") = R.permute_dims(
                    lv3, axes=[0, 2, 1, 3]
                )
                gv: R.Tensor((32, 8, 128, 64), dtype="float32") = lv4
                R.output(gv)
            return gv

    @I.ir_module
    class Expected2:
        @R.function
        def main(
            inp_0: R.Tensor((32, 8, 128, 64), dtype="float32"),
            inp_1: R.Tensor((32, 8, 128, 64), dtype="float32"),
            inp_2: R.Tensor((32, 8, 128, 64), dtype="float32"),
            inp_3: R.Tensor((32, 8, 128, 128), dtype="float32"),
        ) -> R.Tensor((32, 8, 128, 64), dtype="float32"):
            with R.dataflow():
                lv: R.Tensor((32, 128, 8, 64), dtype="float32") = R.permute_dims(
                    inp_0, axes=[0, 2, 1, 3]
                )
                lv1: R.Tensor((32, 128, 8, 64), dtype="float32") = R.permute_dims(
                    inp_1, axes=[0, 2, 1, 3]
                )
                lv2: R.Tensor((32, 128, 8, 64), dtype="float32") = R.permute_dims(
                    inp_2, axes=[0, 2, 1, 3]
                )
                lv3: R.Tensor((32, 128, 8, 64), dtype="float32") = R.nn.attention(
                    lv, lv1, lv2, inp_3, scale=None
                )
                lv4: R.Tensor((32, 8, 128, 64), dtype="float32") = R.permute_dims(
                    lv3, axes=[0, 2, 1, 3]
                )
                gv: R.Tensor((32, 8, 128, 64), dtype="float32") = lv4
                R.output(gv)
            return gv

    @I.ir_module
    class Expected3:
        @R.function
        def main(
            inp_0: R.Tensor((32, 8, 128, 64), dtype="float32"),
            inp_1: R.Tensor((32, 8, 128, 64), dtype="float32"),
            inp_2: R.Tensor((32, 8, 128, 64), dtype="float32"),
        ) -> R.Tensor((32, 8, 128, 64), dtype="float32"):
            with R.dataflow():
                lv: R.Tensor((32, 128, 8, 64), dtype="float32") = R.permute_dims(
                    inp_0, axes=[0, 2, 1, 3]
                )
                lv1: R.Tensor((32, 128, 8, 64), dtype="float32") = R.permute_dims(
                    inp_1, axes=[0, 2, 1, 3]
                )
                lv2: R.Tensor((32, 128, 8, 64), dtype="float32") = R.permute_dims(
                    inp_2, axes=[0, 2, 1, 3]
                )
                lv3: R.Tensor((32, 128, 8, 64), dtype="float32") = R.nn.attention(
                    lv, lv1, lv2, scale=None, causal_mask="TopLeft"
                )
                lv4: R.Tensor((32, 8, 128, 64), dtype="float32") = R.permute_dims(
                    lv3, axes=[0, 2, 1, 3]
                )
                gv: R.Tensor((32, 8, 128, 64), dtype="float32") = lv4
                R.output(gv)
            return gv

    verify_model(
        lambda q, k, v: F.scaled_dot_product_attention(q, k, v),
        [
            ([32, 8, 128, 64], "float32"),
            ([32, 8, 128, 64], "float32"),
            ([32, 8, 128, 64], "float32"),
        ],
        {},
        Expected1,
    )

    verify_model(
        lambda q, k, v, mask: F.scaled_dot_product_attention(q, k, v, mask),
        [
            ([32, 8, 128, 64], "float32"),
            ([32, 8, 128, 64], "float32"),
            ([32, 8, 128, 64], "float32"),
            ([32, 8, 128, 128], "float32"),
        ],
        {},
        Expected2,
    )

    verify_model(
        lambda q, k, v: F.scaled_dot_product_attention(q, k, v, is_causal=True),
        [
            ([32, 8, 128, 64], "float32"),
            ([32, 8, 128, 64], "float32"),
            ([32, 8, 128, 64], "float32"),
        ],
        {},
        Expected3,
    )


def test_sym_size_int():
    class SymSizeInt1(Module):
        def __init__(self, dim):
            super().__init__()
            self.dim = dim

        def forward(self, x):
            return torch.ops.aten.sym_size.int(x, self.dim)

    @I.ir_module
    class Expected1:
        @R.function
        def main(
            inp_0: R.Tensor((1, 3, 4), dtype="float32"),
        ) -> R.Tensor((), dtype="int32"):
            with R.dataflow():
                lv: R.Tensor((), dtype="int32") = R.const(3, "int32")
                gv: R.Tensor((), dtype="int32") = lv
                R.output(gv)
            return gv

    verify_model(SymSizeInt1(dim=1), [([1, 3, 4], "float32")], {}, Expected1)
    verify_model(SymSizeInt1(dim=-2), [([1, 3, 4], "float32")], {}, Expected1)


def test_stack():
    input_info = [
        ([1, 3, 10, 10], "float32"),
        ([1, 3, 10, 10], "float32"),
        ([1, 3, 10, 10], "float32"),
    ]

    class Stack(Module):
        def forward(self, data, data1, data2):
            return torch.stack((data, data1, data2), dim=0)

    @tvm.script.ir_module
    class expected:
        @R.function
        def main(
            inp_0: R.Tensor((1, 3, 10, 10), dtype="float32"),
            inp_1: R.Tensor((1, 3, 10, 10), dtype="float32"),
            inp_2: R.Tensor((1, 3, 10, 10), dtype="float32"),
        ) -> R.Tensor((3, 1, 3, 10, 10), dtype="float32"):
            with R.dataflow():
                lv: R.Tensor((3, 1, 3, 10, 10), dtype="float32") = R.stack(
                    (inp_0, inp_1, inp_2), axis=0
                )
                gv: R.Tensor((3, 1, 3, 10, 10), dtype="float32") = lv
                R.output(gv)
            return gv

    verify_model(Stack(), input_info, {}, expected)


def test_scatter():
    input_info = [([20, 20], "float32"), ([2, 5], "int64"), ([2, 5], "float32")]

    class Scatter(Module):
        def forward(self, data, index, src):
            return data.scatter(dim=0, index=index, src=src)

    @tvm.script.ir_module
    class expected:
        @R.function
        def main(
            inp_0: R.Tensor((20, 20), dtype="float32"),
            inp_1: R.Tensor((2, 5), dtype="int64"),
            inp_2: R.Tensor((2, 5), dtype="float32"),
        ) -> R.Tensor((20, 20), dtype="float32"):
            with R.dataflow():
                lv: R.Tensor((20, 20), dtype="float32") = R.scatter_elements(
                    inp_0, inp_1, inp_2, axis=0, reduction="update"
                )
                gv: R.Tensor((20, 20), dtype="float32") = lv
                R.output(gv)
            return gv

    verify_model(Scatter(), input_info, {}, expected)


def test_slice_scatter():
    class SliceScatter1(Module):
        def forward(self, input, src):
            return torch.slice_scatter(input, src, dim=1, start=1, end=7, step=2)

    @tvm.script.ir_module
    class expected1:
        @R.function
        def main(
            a: R.Tensor((8, 8, 10, 10), dtype="float32"),
            b: R.Tensor((8, 3, 10, 10), dtype="float32"),
        ) -> R.Tensor((8, 8, 10, 10), dtype="float32"):
            with R.dataflow():
                lv: R.Tensor((8, 8, 10, 10), dtype="float32") = R.slice_scatter(
                    a, b, R.prim_value(1), R.prim_value(7), R.prim_value(2), axis=1
                )
                gv: R.Tensor((8, 8, 10, 10), dtype="float32") = lv
                R.output(gv)
            return gv

    class SliceScatter2(Module):
        def forward(self, input, src):
            return torch.slice_scatter(input, src, dim=0, start=0, end=6, step=1)

    @I.ir_module
    class expected2:
        @R.function
        def main(
            a: R.Tensor((8, 16), dtype="float32"), b: R.Tensor((6, 16), dtype="float32")
        ) -> R.Tensor((8, 16), dtype="float32"):
            with R.dataflow():
                lv: R.Tensor((8, 16), dtype="float32") = R.slice_scatter(
                    a, b, R.prim_value(0), R.prim_value(6), R.prim_value(1), axis=0
                )
                gv: R.Tensor((8, 16), dtype="float32") = lv
                R.output(gv)
            return gv

    class SliceScatterNegative(Module):
        def forward(self, input, src):
            return torch.slice_scatter(input, src, dim=1, start=0, end=-2, step=1)

    @tvm.script.ir_module
    class expected_slice_scatter:
        @R.function
        def main(
            a: R.Tensor((2, 5), dtype="float32"), b: R.Tensor((2, 3), dtype="float32")
        ) -> R.Tensor((2, 5), dtype="float32"):
            with R.dataflow():
                lv: R.Tensor((2, 5), dtype="float32") = R.slice_scatter(
                    a, b, R.prim_value(0), R.prim_value(3), R.prim_value(1), axis=1
                )
                gv: R.Tensor((2, 5), dtype="float32") = lv
                R.output(gv)
            return gv

    verify_model(
        SliceScatter1(), [((8, 8, 10, 10), "float32"), ((8, 3, 10, 10), "float32")], {}, expected1
    )
    verify_model(SliceScatter2(), [((8, 16), "float32"), ((6, 16), "float32")], {}, expected2)
    verify_model(
        SliceScatterNegative(),
        [((2, 5), "float32"), ((2, 3), "float32")],
        {},
        expected_slice_scatter,
    )


def test_masked_scatter():
    class MaskedScatter1(Module):
        def forward(self, data, mask, src):
            return data.masked_scatter(mask, src)

    class MaskedScatter2(Module):
        def forward(self, data, mask, src):
            return data.masked_scatter(mask, src)

    @tvm.script.ir_module
    class expected1:
        @R.function
        def main(
            inp_0: R.Tensor((5,), dtype="float32"),
            inp_1: R.Tensor((5,), dtype="bool"),
            inp_2: R.Tensor((10,), dtype="float32"),
        ) -> R.Tensor((5,), dtype="float32"):
            with R.dataflow():
                lv: R.Tensor((5,), dtype="int32") = R.cumsum(
                    inp_1, axis=0, dtype="int32", exclusive=False
                )
                lv1: R.Tensor((5,), dtype="int32") = R.subtract(lv, R.const(1, "int32"))
                lv2: R.Tensor((5,), dtype="float32") = R.take(inp_2, lv1, axis=0)
                lv3: R.Tensor((5,), dtype="float32") = R.where(inp_1, lv2, inp_0)
                gv: R.Tensor((5,), dtype="float32") = lv3
                R.output(gv)
            return gv

    @tvm.script.ir_module
    class expected2:
        @R.function
        def main(
            inp_0: R.Tensor((2, 5), dtype="float32"),
            inp_1: R.Tensor((2, 5), dtype="bool"),
            inp_2: R.Tensor((3, 5), dtype="float32"),
        ) -> R.Tensor((2, 5), dtype="float32"):
            with R.dataflow():
                lv: R.Tensor((10,), dtype="bool") = R.reshape(inp_1, R.shape([10]))
                lv1: R.Tensor((10,), dtype="int32") = R.cumsum(
                    lv, axis=0, dtype="int32", exclusive=False
                )
                lv2: R.Tensor((10,), dtype="int32") = R.subtract(lv1, R.const(1, "int32"))
                lv3: R.Tensor((15,), dtype="float32") = R.reshape(inp_2, R.shape([15]))
                lv4: R.Tensor((10,), dtype="float32") = R.take(lv3, lv2, axis=0)
                lv5: R.Tensor((2, 5), dtype="float32") = R.reshape(lv4, R.shape([2, 5]))
                lv6: R.Tensor((2, 5), dtype="float32") = R.where(inp_1, lv5, inp_0)
                gv: R.Tensor((2, 5), dtype="float32") = lv6
                R.output(gv)
            return gv

    verify_model(
        MaskedScatter1(), [([5], "float32"), ([5], "bool"), ([10], "float32")], {}, expected1
    )
    verify_model(
        MaskedScatter2(),
        [([2, 5], "float32"), ([2, 5], "bool"), ([3, 5], "float32")],
        {},
        expected2,
    )


def test_is_floating_point():
    class IsFloatingPoint(Module):
        def forward(self, x):
            return torch.is_floating_point(x)

    @tvm.script.ir_module
    class Expected:
        @R.function
        def main(inp_0: R.Tensor((2, 3), dtype="float32")) -> R.Tensor((), dtype="bool"):
            with R.dataflow():
                gv: R.Tensor((), dtype="bool") = R.const(True, "bool")
                R.output(gv)
            return gv

    verify_model(IsFloatingPoint(), [([2, 3], "float32")], {}, Expected)


def test_gather():
    class Gather0(Module):
        def forward(self, data, indices):
            return torch.gather(data, 0, indices)

    class Gather1(Module):
        def forward(self, data, indices):
            return torch.gather(data, 1, indices)

    class Gather2(Module):
        def forward(self, data, indices):
            return torch.gather(data, -1, indices)

    class Gather3(Module):
        def forward(self, data, indices):
            return torch.gather(data, -2, indices)

    @tvm.script.ir_module
    class Expected0:
        @R.function
        def main(
            inp_0: R.Tensor((2, 3), dtype="float32"),
            inp_1: R.Tensor((2, 3), dtype="int32"),
        ) -> R.Tensor((2, 3), dtype="float32"):
            with R.dataflow():
                lv: R.Tensor((2, 3), dtype="float32") = R.gather_elements(inp_0, inp_1, axis=0)
                gv: R.Tensor((2, 3), dtype="float32") = lv
                R.output(gv)
            return gv

    @tvm.script.ir_module
    class Expected1:
        @R.function
        def main(
            inp_0: R.Tensor((2, 3), dtype="float32"),
            inp_1: R.Tensor((2, 3), dtype="int32"),
        ) -> R.Tensor((2, 3), dtype="float32"):
            with R.dataflow():
                lv: R.Tensor((2, 3), dtype="float32") = R.gather_elements(inp_0, inp_1, axis=1)
                gv: R.Tensor((2, 3), dtype="float32") = lv
                R.output(gv)
            return gv

    @tvm.script.ir_module
    class Expected2:
        @R.function
        def main(
            inp_0: R.Tensor((2, 3), dtype="float32"),
            inp_1: R.Tensor((2, 3), dtype="int32"),
        ) -> R.Tensor((2, 3), dtype="float32"):
            with R.dataflow():
                lv: R.Tensor((2, 3), dtype="float32") = R.gather_elements(inp_0, inp_1, axis=-1)
                gv: R.Tensor((2, 3), dtype="float32") = lv
                R.output(gv)
            return gv

    @tvm.script.ir_module
    class Expected3:
        @R.function
        def main(
            inp_0: R.Tensor((2, 3), dtype="float32"),
            inp_1: R.Tensor((2, 3), dtype="int32"),
        ) -> R.Tensor((2, 3), dtype="float32"):
            with R.dataflow():
                lv: R.Tensor((2, 3), dtype="float32") = R.gather_elements(inp_0, inp_1, axis=-2)
                gv: R.Tensor((2, 3), dtype="float32") = lv
                R.output(gv)
            return gv

    verify_model(Gather0(), [([2, 3], "float32"), ([2, 3], "int32")], {}, Expected0)
    verify_model(Gather1(), [([2, 3], "float32"), ([2, 3], "int32")], {}, Expected1)
    verify_model(Gather2(), [([2, 3], "float32"), ([2, 3], "int32")], {}, Expected2)
    verify_model(Gather3(), [([2, 3], "float32"), ([2, 3], "int32")], {}, Expected3)


def test_index_put():
    # Test case 1: 1D input
    class IndexPut1D(Module):
        def forward(self, data, indices_0, values):
            indices_tuple = (indices_0,)
            return data.index_put_(indices_tuple, values, accumulate=False)

    input_info_1d = [((64,), "float32"), ((128,), "int64"), ((128,), "float32")]

    @I.ir_module
    class Expected1D:
        @R.function
        def main(
            data: R.Tensor((64,), dtype="float32"),
            indices_0: R.Tensor((128,), dtype="int64"),
            values: R.Tensor((128,), dtype="float32"),
        ) -> R.Tensor((64,), dtype="float32"):
            with R.dataflow():
                lv: R.Tensor((64,), dtype="float32") = R.index_put(
                    data, R.tuple(indices_0), values, accumulate=False
                )
                gv: R.Tensor((64,), dtype="float32") = lv
                R.output(gv)
            return gv

    # Test case 2: 2D input
    class IndexPut2D(Module):
        def forward(self, data, indices_0, indices_1, values):
            indices_tuple = (indices_0, indices_1)
            return data.index_put_(indices_tuple, values, accumulate=False)

    input_info_2d = [
        ((32, 64), "float32"),
        ((128,), "int64"),
        ((128,), "int64"),
        ((128,), "float32"),
    ]

    @I.ir_module
    class Expected2D:
        @R.function
        def main(
            data: R.Tensor((32, 64), dtype="float32"),
            indices_0: R.Tensor((128,), dtype="int64"),
            indices_1: R.Tensor((128,), dtype="int64"),
            values: R.Tensor((128,), dtype="float32"),
        ) -> R.Tensor((32, 64), dtype="float32"):
            with R.dataflow():
                lv: R.Tensor((32, 64), dtype="float32") = R.index_put(
                    data, R.tuple(indices_0, indices_1), values, accumulate=False
                )
                gv: R.Tensor((32, 64), dtype="float32") = lv
                R.output(gv)
            return gv

    # Test case 3: 3D input
    class IndexPut3D(Module):
        def forward(self, data, indices_0, indices_1, indices_2, values):
            indices_tuple = (indices_0, indices_1, indices_2)
            return data.index_put_(indices_tuple, values, accumulate=False)

    input_info_3d = [
        ((16, 32, 64), "float32"),
        ((128,), "int64"),
        ((128,), "int64"),
        ((128,), "int64"),
        ((128,), "float32"),
    ]

    @I.ir_module
    class Expected3D:
        @R.function
        def main(
            data: R.Tensor((16, 32, 64), dtype="float32"),
            indices_0: R.Tensor((128,), dtype="int64"),
            indices_1: R.Tensor((128,), dtype="int64"),
            indices_2: R.Tensor((128,), dtype="int64"),
            values: R.Tensor((128,), dtype="float32"),
        ) -> R.Tensor((16, 32, 64), dtype="float32"):
            with R.dataflow():
                lv: R.Tensor((16, 32, 64), dtype="float32") = R.index_put(
                    data, R.tuple(indices_0, indices_1, indices_2), values, accumulate=False
                )
                gv: R.Tensor((16, 32, 64), dtype="float32") = lv
                R.output(gv)
            return gv

    # Test case 4: 4D input
    class IndexPut4D(Module):
        def forward(self, data, indices_0, indices_1, indices_2, indices_3, values):
            indices_tuple = (indices_0, indices_1, indices_2, indices_3)
            return data.index_put_(indices_tuple, values, accumulate=False)

    input_info_4d = [
        ((8, 16, 32, 64), "float32"),
        ((128,), "int64"),
        ((128,), "int64"),
        ((128,), "int64"),
        ((128,), "int64"),
        ((128,), "float32"),
    ]

    @I.ir_module
    class Expected4D:
        @R.function
        def main(
            data: R.Tensor((8, 16, 32, 64), dtype="float32"),
            indices_0: R.Tensor((128,), dtype="int64"),
            indices_1: R.Tensor((128,), dtype="int64"),
            indices_2: R.Tensor((128,), dtype="int64"),
            indices_3: R.Tensor((128,), dtype="int64"),
            values: R.Tensor((128,), dtype="float32"),
        ) -> R.Tensor((8, 16, 32, 64), dtype="float32"):
            with R.dataflow():
                lv: R.Tensor((8, 16, 32, 64), dtype="float32") = R.index_put(
                    data,
                    R.tuple(indices_0, indices_1, indices_2, indices_3),
                    values,
                    accumulate=False,
                )
                gv: R.Tensor((8, 16, 32, 64), dtype="float32") = lv
                R.output(gv)
            return gv

    # Test case 5: 5D input
    class IndexPut5D(Module):
        def forward(self, data, indices_0, indices_1, indices_2, indices_3, indices_4, values):
            indices_tuple = (indices_0, indices_1, indices_2, indices_3, indices_4)
            return data.index_put_(indices_tuple, values, accumulate=False)

    input_info_5d = [
        ((4, 8, 16, 32, 64), "float32"),
        ((128,), "int64"),
        ((128,), "int64"),
        ((128,), "int64"),
        ((128,), "int64"),
        ((128,), "int64"),
        ((128,), "float32"),
    ]

    @I.ir_module
    class Expected5D:
        @R.function
        def main(
            data: R.Tensor((4, 8, 16, 32, 64), dtype="float32"),
            indices_0: R.Tensor((128,), dtype="int64"),
            indices_1: R.Tensor((128,), dtype="int64"),
            indices_2: R.Tensor((128,), dtype="int64"),
            indices_3: R.Tensor((128,), dtype="int64"),
            indices_4: R.Tensor((128,), dtype="int64"),
            values: R.Tensor((128,), dtype="float32"),
        ) -> R.Tensor((4, 8, 16, 32, 64), dtype="float32"):
            with R.dataflow():
                lv: R.Tensor((4, 8, 16, 32, 64), dtype="float32") = R.index_put(
                    data,
                    R.tuple(indices_0, indices_1, indices_2, indices_3, indices_4),
                    values,
                    accumulate=False,
                )
                gv: R.Tensor((4, 8, 16, 32, 64), dtype="float32") = lv
                R.output(gv)
            return gv

    # Run verification for each case
    verify_model(IndexPut1D(), input_info_1d, {}, Expected1D)
    verify_model(IndexPut2D(), input_info_2d, {}, Expected2D)
    verify_model(IndexPut3D(), input_info_3d, {}, Expected3D)
    verify_model(IndexPut4D(), input_info_4d, {}, Expected4D)
    verify_model(IndexPut5D(), input_info_5d, {}, Expected5D)


def test_flip():
    class Flip0(Module):
        def forward(self, data):
            return torch.flip(data, [0])

    class Flip1(Module):
        def forward(self, data):
            return torch.flip(data, [1])

    @tvm.script.ir_module
    class Expected0:
        @R.function
        def main(
            inp_0: R.Tensor((2, 2), dtype="float32"),
        ) -> R.Tensor((2, 2), dtype="float32"):
            with R.dataflow():
                lv: R.Tensor((2, 2), dtype="float32") = R.flip(inp_0, axis=0)
                gv: R.Tensor((2, 2), dtype="float32") = lv
                R.output(gv)
            return gv

    @tvm.script.ir_module
    class Expected1:
        @R.function
        def main(
            inp_0: R.Tensor((2, 2), dtype="float32"),
        ) -> R.Tensor((2, 2), dtype="float32"):
            with R.dataflow():
                lv: R.Tensor((2, 2), dtype="float32") = R.flip(inp_0, axis=1)
                gv: R.Tensor((2, 2), dtype="float32") = lv
                R.output(gv)
            return gv

    verify_model(Flip0(), [([2, 2], "float32")], {}, Expected0)
    verify_model(Flip1(), [([2, 2], "float32")], {}, Expected1)


def test_flip_multi_axis():
    class FlipMulti(Module):
        def forward(self, data):
            return torch.flip(data, [0, 1])

    @tvm.script.ir_module
    class ExpectedMulti:
        @R.function
        def main(
            inp_0: R.Tensor((2, 3), dtype="float32"),
        ) -> R.Tensor((2, 3), dtype="float32"):
            with R.dataflow():
                lv: R.Tensor((2, 3), dtype="float32") = R.flip(inp_0, axis=0)
                lv1: R.Tensor((2, 3), dtype="float32") = R.flip(lv, axis=1)
                gv: R.Tensor((2, 3), dtype="float32") = lv1
                R.output(gv)
            return gv

    verify_model(FlipMulti(), [([2, 3], "float32")], {}, ExpectedMulti)


def test_take():
    class Take(Module):
        def forward(self, data, indices):
            return torch.take(data, indices)

    @tvm.script.ir_module
    class Expected:
        @R.function
        def main(
            inp_0: R.Tensor((5,), dtype="float32"),
            inp_1: R.Tensor((3,), dtype="int32"),
        ) -> R.Tensor((3,), dtype="float32"):
            with R.dataflow():
                lv: R.Tensor((3,), dtype="int32") = R.astype(inp_1, "int32")
                lv1: R.Tensor((3,), dtype="float32") = R.take(inp_0, lv)
                gv: R.Tensor((3,), dtype="float32") = lv1
                R.output(gv)
            return gv

    verify_model(Take(), [([5], "float32"), ([3], "int32")], {}, Expected)


def test_one_hot():
    class OneHot(Module):
        def forward(self, indices):
            return torch.nn.functional.one_hot(indices, num_classes=10)

    @tvm.script.ir_module
    class Expected:
        @R.function
        def main(
            inp_0: R.Tensor((5,), dtype="int32"),
        ) -> R.Tensor((5, 10), dtype="int64"):
            with R.dataflow():
                lv: R.Tensor((5, 10), dtype="int64") = R.one_hot(
                    inp_0, R.prim_value(1), R.prim_value(0), depth=10, axis=-1
                )
                gv: R.Tensor((5, 10), dtype="int64") = lv
                R.output(gv)

            return gv

    verify_model(OneHot(), [([5], "int32")], {}, Expected)


def test_empty_like():
    class EmptyLike(Module):
        def forward(self, data):
            return torch.empty_like(data)

    @tvm.script.ir_module
    class Expected:
        @R.function
        def main(
            inp_0: R.Tensor((5,), dtype="float32"),
        ) -> R.Tensor((5,), dtype="float32"):
            with R.dataflow():
                lv: R.Tensor((5,), dtype="float32") = R.zeros_like(inp_0)
                gv: R.Tensor((5,), dtype="float32") = lv
                R.output(gv)
            return gv

    verify_model(EmptyLike(), [([5], "float32")], {}, Expected)


def test_ones_like():
    class OnesLike(Module):
        def forward(self, data):
            return torch.ones_like(data)

    @tvm.script.ir_module
    class Expected:
        @R.function
        def main(inp_0: R.Tensor((128, 128), dtype="float32")) -> R.Tensor(
            (128, 128), dtype="float32"
        ):
            with R.dataflow():
                lv: R.Tensor((128, 128), dtype="float32") = R.ones_like(inp_0)
                gv: R.Tensor((128, 128), dtype="float32") = lv
                R.output(gv)
            return gv

    verify_model(OnesLike(), [([128, 128], "float32")], {}, Expected)


def test_zero_inplace():
    class ZeroInplace(Module):
        def forward(self, data):
            return data.zero_()

    @tvm.script.ir_module
    class Expected:
        @R.function
        def main(inp_0: R.Tensor((128, 128), dtype="float32")) -> R.Tensor(
            (128, 128), dtype="float32"
        ):
            with R.dataflow():
                lv: R.Tensor((128, 128), dtype="float32") = R.zeros_like(inp_0)
                gv: R.Tensor((128, 128), dtype="float32") = lv
                R.output(gv)
            return gv

    verify_model(ZeroInplace(), [([128, 128], "float32")], {}, Expected)


def test_zeros_like():
    class ZerosLike(Module):
        def forward(self, data):
            return torch.zeros_like(data)

    @tvm.script.ir_module
    class Expected:
        @R.function
        def main(inp_0: R.Tensor((128, 128), dtype="float32")) -> R.Tensor(
            (128, 128), dtype="float32"
        ):
            with R.dataflow():
                lv: R.Tensor((128, 128), dtype="float32") = R.zeros_like(inp_0)
                gv: R.Tensor((128, 128), dtype="float32") = lv
                R.output(gv)
            return gv

    verify_model(ZerosLike(), [([128, 128], "float32")], {}, Expected)


def test_type_as():
    class TypeAs(Module):
        def forward(self, data, other):
            return data.type_as(other)

    @tvm.script.ir_module
    class Expected:
        @R.function
        def main(
            inp_0: R.Tensor((128, 128), dtype="float16"),
            inp_1: R.Tensor((128, 128), dtype="float32"),
        ) -> R.Tensor((128, 128), dtype="float32"):
            with R.dataflow():
                lv: R.Tensor((128, 128), dtype="float32") = R.astype(inp_0, dtype="float32")
                gv: R.Tensor((128, 128), dtype="float32") = lv
                R.output(gv)
            return gv

    verify_model(TypeAs(), [([128, 128], "float16"), ([128, 128], "float32")], {}, Expected)


def test_item():
    class Item(Module):
        def forward(self, data):
            return data.item()

    @tvm.script.ir_module
    class Expected:
        @R.function
        def main(inp_0: R.Tensor((1,), dtype="float32")) -> R.Tensor((), dtype="float32"):
            with R.dataflow():
                lv: R.Tensor((), dtype="float32") = R.take(inp_0, R.const(0, "int64"), axis=0)
                gv: R.Tensor((), dtype="float32") = lv
                R.output(gv)
            return gv

    verify_model(
        Item(),
        [
            (
                [1],
                "float32",
            )
        ],
        {},
        Expected,
    )


def test_numel():
    class Numel(Module):
        def forward(self, data):
            return torch.numel(data)

    @tvm.script.ir_module
    class Expected:
        @R.function
        def main(
            inp_0: R.Tensor((5, 3), dtype="float32"),
        ) -> R.Tensor((), dtype="int32"):
            with R.dataflow():
                gv: R.Tensor((), dtype="int32") = R.const(15, "int32")
                R.output(gv)
            return gv

    verify_model(Numel(), [([5, 3], "float32")], {}, Expected)


def test_select():
    class Select(Module):
        def forward(self, data):
            return torch.select(data, 0, 1)

    @tvm.script.ir_module
    class Expected:
        @R.function
        def main(
            inp_0: R.Tensor((5, 3), dtype="float32"),
        ) -> R.Tensor((3,), dtype="float32"):
            with R.dataflow():
                lv: R.Tensor((3,), dtype="float32") = R.take(inp_0, R.const(1, "int64"), axis=0)
                gv: R.Tensor((3,), dtype="float32") = lv
                R.output(gv)
            return gv

    verify_model(Select(), [([5, 3], "float32")], {}, Expected)


def test_inplace_copy():
    class Inplace_Copy(Module):
        def forward(self, x, y):
            x.copy_(y)
            return x

    @tvm.script.ir_module
    class Expected:
        @R.function
        def main(
            inp_0: R.Tensor((1, 2, 3, 4), dtype="float32"),
            inp_1: R.Tensor((1, 2, 3, 4), dtype="float32"),
        ) -> R.Tensor((1, 2, 3, 4), dtype="float32"):
            with R.dataflow():
                lv: R.Tensor((1, 2, 3, 4), dtype="float32") = R.broadcast_to(
                    inp_1, R.shape([1, 2, 3, 4])
                )
                gv: R.Tensor((1, 2, 3, 4), dtype="float32") = lv
                R.output(gv)
            return gv

    class CopyBroadcast(Module):
        def forward(self, x, src):
            x.copy_(src)
            return x

    @tvm.script.ir_module
    class expected_copy:
        @R.function
        def main(
            x: R.Tensor((2, 3), dtype="float32"), src: R.Tensor((), dtype="int64")
        ) -> R.Tensor((2, 3), dtype="float32"):
            with R.dataflow():
                lv: R.Tensor((), dtype="float32") = R.astype(src, dtype="float32")
                lv1: R.Tensor((2, 3), dtype="float32") = R.broadcast_to(lv, (2, 3))
                gv: R.Tensor((2, 3), dtype="float32") = lv1
                R.output(gv)
            return gv

    verify_model(
        Inplace_Copy(),
        [((1, 2, 3, 4), "float32"), ((1, 2, 3, 4), "float32")],
        {},
        Expected,
    )
    verify_model(CopyBroadcast(), [((2, 3), "float32"), ((), "int64")], {}, expected_copy)


def test_clone():
    class Clone(Module):
        def forward(self, x):
            return x.clone()

    @tvm.script.ir_module
    class Expected:
        @R.function
        def main(
            inp_0: R.Tensor((5, 3), dtype="float32"),
        ) -> R.Tensor((5, 3), dtype="float32"):
            with R.dataflow():
                gv: R.Tensor((5, 3), dtype="float32") = inp_0
                R.output(gv)
            return gv

    verify_model(Clone(), [([5, 3], "float32")], {}, Expected)


def test_lerp():
    class Lerp(Module):
        def forward(self, start, end, weight):
            return torch.lerp(start, end, weight)

    @tvm.script.ir_module
    class Expected:
        @R.function
        def main(
            inp_0: R.Tensor((5, 3), dtype="float32"),
            inp_1: R.Tensor((5, 3), dtype="float32"),
            inp_2: R.Tensor((5, 3), dtype="float32"),
        ) -> R.Tensor((5, 3), dtype="float32"):
            with R.dataflow():
                lv: R.Tensor((5, 3), dtype="float32") = R.add(
                    inp_0, R.multiply(inp_2, R.subtract(inp_1, inp_0))
                )
                gv: R.Tensor((5, 3), dtype="float32") = lv
                R.output(gv)
            return gv

    verify_model(
        Lerp(), [([5, 3], "float32"), ([5, 3], "float32"), ([5, 3], "float32")], {}, Expected
    )


def test_std():
    class Std(Module):
        def forward(self, x):
            return torch.std(x)

    @tvm.script.ir_module
    class Expected:
        @R.function
        def main(
            inp_0: R.Tensor((5, 3), dtype="float32"),
        ) -> R.Tensor((), dtype="float32"):
            with R.dataflow():
                lv: R.Tensor((), dtype="float32") = R.std(inp_0, axis=None, keepdims=False)
                gv: R.Tensor((), dtype="float32") = lv
                R.output(gv)
            return gv

    verify_model(Std(), [([5, 3], "float32")], {}, Expected)


def test_var():
    class Var(Module):
        def forward(self, x):
            return torch.var(x)

    @tvm.script.ir_module
    class Expected:
        @R.function
        def main(
            inp_0: R.Tensor((5, 3), dtype="float32"),
        ) -> R.Tensor((), dtype="float32"):
            with R.dataflow():
                lv: R.Tensor((), dtype="float32") = R.variance(inp_0, axis=None, keepdims=False)
                gv: R.Tensor((), dtype="float32") = lv
                R.output(gv)
            return gv

    verify_model(Var(), [([5, 3], "float32")], {}, Expected)


@pytest.mark.parametrize(
    "torch_dtype,relax_dtype",
    [(torch.float32, "float32"), (torch.bool, "bool")],
)
def test_prod(torch_dtype, relax_dtype):
    class Prod(Module):
        def forward(self, x):
            return torch.prod(x, dtype=torch_dtype)

    @tvm.script.ir_module
    class Expected:
        @R.function
        def main(
            inp_0: R.Tensor((5, 3), dtype=relax_dtype),
        ) -> R.Tensor((), dtype=relax_dtype):
            with R.dataflow():
                lv: R.Tensor((), dtype=relax_dtype) = R.prod(inp_0, axis=None, keepdims=False)
                gv: R.Tensor((), dtype=relax_dtype) = lv
                R.output(gv)
            return gv

    verify_model(Prod(), [([5, 3], relax_dtype)], {}, Expected)


def test_cumprod():
    class Cumprod(Module):
        def forward(self, x):
            return torch.cumprod(x, 0)

    @tvm.script.ir_module
    class Expected:
        @R.function
        def main(
            inp_0: R.Tensor((5, 3), dtype="float32"),
        ) -> R.Tensor((5, 3), dtype="float32"):
            with R.dataflow():
                lv: R.Tensor((5, 3), dtype="float32") = R.cumprod(inp_0, axis=0, exclusive=False)
                gv: R.Tensor((5, 3), dtype="float32") = lv
                R.output(gv)
            return gv

    verify_model(Cumprod(), [([5, 3], "float32")], {}, Expected)


def test_where():
    class Where(Module):
        def forward(self, condition, x, y):
            return torch.where(condition, x, y)

    @tvm.script.ir_module
    class Expected:
        @R.function
        def main(
            inp_0: R.Tensor((5, 3), dtype="bool"),
            inp_1: R.Tensor((5, 3), dtype="float32"),
            inp_2: R.Tensor((5, 3), dtype="float32"),
        ) -> R.Tensor((5, 3), dtype="float32"):
            with R.dataflow():
                lv: R.Tensor((5, 3), dtype="float32") = R.where(inp_0, inp_1, inp_2)
                gv: R.Tensor((5, 3), dtype="float32") = lv
                R.output(gv)
            return gv

    verify_model(
        Where(), [([5, 3], "bool"), ([5, 3], "float32"), ([5, 3], "float32")], {}, Expected
    )


def test_bucketize():
    class Bucketize(Module):
        def forward(self, input_tensor, boundaries):
            return torch.bucketize(input_tensor, boundaries)

    @tvm.script.ir_module
    class Expected:
        @R.function
        def main(
            input: R.Tensor((5, 3), dtype="float32"), boundaries: R.Tensor((10,), dtype="float32")
        ) -> R.Tensor((5, 3), dtype="int64"):
            with R.dataflow():
                lv: R.Tensor((5, 3), dtype="int64") = R.bucketize(
                    input, boundaries, out_int32=False, right=False
                )
                gv: R.Tensor((5, 3), dtype="int64") = lv
                R.output(gv)
            return gv

    verify_model(Bucketize(), [([5, 3], "float32"), ([10], "float32")], {}, Expected)


def test_argsort():
    class Argsort(Module):
        def forward(self, x):
            return torch.argsort(x, dim=1, descending=True)

    @tvm.script.ir_module
    class Expected:
        @R.function
        def main(
            inp_0: R.Tensor((5, 3), dtype="float32"),
        ) -> R.Tensor((5, 3), dtype="int32"):
            with R.dataflow():
                lv: R.Tensor((5, 3), dtype="int32") = R.argsort(inp_0, axis=1, descending=True)
                gv: R.Tensor((5, 3), dtype="int32") = lv
                R.output(gv)
            return gv

    verify_model(Argsort(), [([5, 3], "float32")], {}, Expected)


def test_sort():
    class Sort(Module):
        def forward(self, x):
            return torch.sort(x, dim=1, descending=True)

    @tvm.script.ir_module
    class Expected:
        @R.function
        def main(inp_0: R.Tensor((5, 3), dtype="float32")) -> R.Tuple(
            R.Tensor((5, 3), dtype="float32"), R.Tensor((5, 3), dtype="int32")
        ):
            with R.dataflow():
                lv: R.Tensor((5, 3), dtype="int32") = R.argsort(
                    inp_0, axis=1, descending=True, dtype="int32"
                )
                lv1: R.Tensor((5, 3), dtype="float32") = R.gather_elements(inp_0, lv, axis=1)
                lv2: R.Tuple(R.Tensor((5, 3), dtype="float32"), R.Tensor((5, 3), dtype="int32")) = (
                    lv1,
                    lv,
                )
                gv: R.Tuple(R.Tensor((5, 3), dtype="float32"), R.Tensor((5, 3), dtype="int32")) = (
                    lv2
                )
                R.output(gv)
            return gv

    verify_model(Sort(), [([5, 3], "float32")], {}, Expected)


def test_topk():
    class Topk(Module):
        def forward(self, x):
            return torch.topk(x, k=2, dim=1, largest=True, sorted=True)

    @tvm.script.ir_module
    class Expected:
        @R.function
        def main(
            inp_0: R.Tensor((5, 3), dtype="float32"),
        ) -> R.Tuple(R.Tensor((5, 2), dtype="float32"), R.Tensor((5, 2), dtype="int64")):
            with R.dataflow():
                lv: R.Tuple(R.Tensor((5, 2), dtype="float32"), R.Tensor((5, 2), dtype="int64")) = (
                    R.topk(inp_0, k=2, axis=1, ret_type="both", largest=True, dtype="int64")
                )
                gv: R.Tuple(R.Tensor((5, 2), dtype="float32"), R.Tensor((5, 2), dtype="int64")) = lv
                R.output(gv)
            return gv

    verify_model(Topk(), [([5, 3], "float32")], {}, Expected)


def test_broadcast_to():
    class BroadcastTo(Module):
        def forward(self, x):
            return torch.broadcast_to(x, (5, 3))

    @tvm.script.ir_module
    class Expected:
        @R.function
        def main(
            inp_0: R.Tensor((5, 1), dtype="float32"),
        ) -> R.Tensor((5, 3), dtype="float32"):
            with R.dataflow():
                lv: R.Tensor((5, 3), dtype="float32") = R.broadcast_to(inp_0, (5, 3))
                gv: R.Tensor((5, 3), dtype="float32") = lv
                R.output(gv)
            return gv

    verify_model(BroadcastTo(), [([5, 1], "float32")], {}, Expected)


def test_narrow():
    class Narrow(Module):
        def forward(self, x):
            return torch.narrow(x, 1, 0, 2)

    @tvm.script.ir_module
    class Expected:
        @R.function
        def main(
            inp_0: R.Tensor((5, 3), dtype="float32"),
        ) -> R.Tensor((5, 2), dtype="float32"):
            with R.dataflow():
                lv: R.Tensor((5, 2), dtype="float32") = R.strided_slice(
                    inp_0, axes=[1], begin=[0], end=[2]
                )
                gv: R.Tensor((5, 2), dtype="float32") = lv
                R.output(gv)

            return gv

    verify_model(Narrow(), [([5, 3], "float32")], {}, Expected)


def test_norm():
    input_info = [([1, 3, 5, 3], "float32")]

    class Norm(Module):
        def __init__(self, p, dim=None, keepdim=False):
            super().__init__()
            self.p = p
            self.dim = dim
            self.keepdim = keepdim

        def forward(self, x):
            return torch.norm(x, p=self.p, dim=self.dim, keepdim=self.keepdim)

    @tvm.script.ir_module
    class Expected1:
        @R.function
        def main(
            inp_0: R.Tensor((1, 3, 5, 3), dtype="float32"),
        ) -> R.Tensor((), dtype="float32"):
            with R.dataflow():
                lv: R.Tensor((), dtype="float32") = R.max(R.abs(inp_0), axis=None, keepdims=False)
                gv: R.Tensor((), dtype="float32") = lv
                R.output(gv)
            return gv

    @tvm.script.ir_module
    class Expected2:
        @R.function
        def main(
            inp_0: R.Tensor((1, 3, 5, 3), dtype="float32"),
        ) -> R.Tensor((), dtype="float32"):
            with R.dataflow():
                lv: R.Tensor((), dtype="float32") = R.min(R.abs(inp_0), axis=None, keepdims=False)
                gv: R.Tensor((), dtype="float32") = lv
                R.output(gv)
            return gv

    @tvm.script.ir_module
    class Expected3:
        @R.function
        def main(
            inp_0: R.Tensor((1, 3, 5, 3), dtype="float32"),
        ) -> R.Tensor((), dtype="float32"):
            with R.dataflow():
                lv: R.Tensor((1, 3, 5, 3), dtype="float32") = R.abs(inp_0)
                lv1: R.Tensor((1, 3, 5, 3), dtype="float32") = R.power(lv, R.const(2, "float32"))
                lv2: R.Tensor((), dtype="float32") = R.sum(lv1, axis=None, keepdims=False)
                lv3: R.Tensor((), dtype="float32") = R.power(lv2, R.const(0.5, "float32"))
                gv: R.Tensor((), dtype="float32") = lv3
                R.output(gv)
            return gv

    @tvm.script.ir_module
    class Expected4:
        @R.function
        def main(
            inp_0: R.Tensor((1, 3, 5, 3), dtype="float32"),
        ) -> R.Tensor((), dtype="float32"):
            with R.dataflow():
                lv: R.Tensor((1, 3, 5, 3), dtype="float32") = R.abs(inp_0)
                lv1: R.Tensor((1, 3, 5, 3), dtype="float32") = R.power(lv, R.const(1.0, "float32"))
                lv2: R.Tensor((), dtype="float32") = R.sum(lv1, axis=None, keepdims=False)
                lv3: R.Tensor((), dtype="float32") = R.power(lv2, R.const(1.0, "float32"))
                gv: R.Tensor((), dtype="float32") = lv3
                R.output(gv)
            return gv

    @tvm.script.ir_module
    class Expected5:
        @R.function
        def main(
            inp_0: R.Tensor((1, 3, 5, 3), dtype="float32"),
        ) -> R.Tensor((), dtype="float32"):
            with R.dataflow():
                lv: R.Tensor((1, 3, 5, 3), dtype="float32") = R.abs(inp_0)
                lv1: R.Tensor((1, 3, 5, 3), dtype="float32") = R.power(lv, R.const(-4, "float32"))
                lv2: R.Tensor((), dtype="float32") = R.sum(lv1, axis=None, keepdims=False)
                lv3: R.Tensor((), dtype="float32") = R.power(lv2, R.const(-0.25, "float32"))
                gv: R.Tensor((), dtype="float32") = lv3
                R.output(gv)
            return gv

    @tvm.script.ir_module
    class Expected6:
        @R.function
        def main(
            inp_0: R.Tensor((1, 3, 5, 3), dtype="float32"),
        ) -> R.Tensor((), dtype="float32"):
            with R.dataflow():
                lv: R.Tensor((1, 3, 5, 3), dtype="float32") = R.abs(inp_0)
                lv1: R.Tensor((1, 3, 5, 3), dtype="float32") = R.power(lv, R.const(0.5, "float32"))
                lv2: R.Tensor((), dtype="float32") = R.sum(lv1, axis=None, keepdims=False)
                lv3: R.Tensor((), dtype="float32") = R.power(lv2, R.const(2, "float32"))
                gv: R.Tensor((), dtype="float32") = lv3
                R.output(gv)
            return gv

    @tvm.script.ir_module
    class Expected7:
        @R.function
        def main(
            inp_0: R.Tensor((1, 3, 5, 3), dtype="float32"),
        ) -> R.Tensor((), dtype="float32"):
            with R.dataflow():
                lv: R.Tensor((1, 3, 5, 3), dtype="float32") = R.multiply(inp_0, inp_0)
                lv1: R.Tensor((), dtype="float32") = R.sum(lv, axis=None, keepdims=False)
                lv2: R.Tensor((), dtype="float32") = R.sqrt(lv1)
                gv: R.Tensor((), dtype="float32") = lv2
                R.output(gv)
            return gv

    norms = [
        ((float("inf"), None, False), Expected1),
        ((float("-inf"), None, False), Expected2),
        ((float(2), None, False), Expected3),
        ((1.0, None, False), Expected4),
        ((float(-4), None, True), Expected5),
        ((0.5, None, True), Expected6),
        (("fro", None, False), Expected7),
    ]

    for (p, dim, keepdim), expected in norms:
        verify_model(Norm(p, dim=dim, keepdim=keepdim), input_info, {}, expected)


@pytest.mark.parametrize(
    "torch_dtype, relax_dtype",
    [
        # Float types
        (torch.float16, "float16"),
        (torch.float32, "float32"),
        (torch.float64, "float64"),
        (torch.bfloat16, "bfloat16"),
        # Signed integer types
        (torch.int8, "int8"),
        (torch.int16, "int16"),
        (torch.int32, "int32"),
        (torch.int64, "int64"),
        # Unsigned integer types
        (torch.uint8, "uint8"),
        (torch.uint16, "uint16"),
        (torch.uint32, "uint32"),
        (torch.uint64, "uint64"),
        # Boolean
        (torch.bool, "bool"),
    ],
)
def test_dtypes(torch_dtype, relax_dtype):
    class Model(Module):
        def forward(self, lhs: torch.Tensor, rhs: torch.Tensor):
            return torch.ops.aten.add(lhs, rhs)

    @tvm.script.ir_module
    class Expected:
        @R.function
        def main(
            lhs: R.Tensor((10, 10), dtype=relax_dtype),
            rhs: R.Tensor((10, 10), dtype=relax_dtype),
        ) -> R.Tensor((10, 10), dtype=relax_dtype):
            with R.dataflow():
                lv: R.Tensor((10, 10), dtype=relax_dtype) = relax.op.add(lhs, rhs)
                gv: R.Tensor((10, 10), dtype=relax_dtype) = lv
                R.output(gv)
            return gv

    verify_model(Model(), [([10, 10], torch_dtype), ([10, 10], torch_dtype)], {}, Expected)


def test_eye():
    import numpy as np

    class Eye(Module):
        def forward(self, input):
            return torch.eye(3)

    graph_model = fx.symbolic_trace(Eye())
    mod = from_fx(graph_model, [([3, 3], "float32")])
    assert len(mod["main"].body.blocks) == 1
    assert len(mod["main"].body.blocks[0].bindings) == 1
    assert isinstance(mod["main"].body.blocks[0].bindings[0].value, relax.Constant)
    tvm.testing.assert_allclose(
        mod["main"].body.blocks[0].bindings[0].value.data.numpy(),
        np.eye(3, dtype="float32"),
    )


def test_linspace():
    import numpy as np

    class Linspace(Module):
        def forward(self, input):
            return torch.linspace(0, 1, steps=9)

    graph_model = fx.symbolic_trace(Linspace())
    mod = from_fx(graph_model, [([9, 9], "float32")])
    assert len(mod["main"].body.blocks) == 1
    assert len(mod["main"].body.blocks[0].bindings) == 1
    assert isinstance(mod["main"].body.blocks[0].bindings[0].value, relax.Constant)
    tvm.testing.assert_allclose(
        mod["main"].body.blocks[0].bindings[0].value.data.numpy(),
        np.linspace(0, 1, num=9, dtype="float32"),
    )


def test_round():
    input_info = [([3, 4], "float32")]

    class Round(Module):
        def __init__(self, decimals=0):
            super().__init__()
            self.decimals = decimals

        def forward(self, x):
            if self.decimals == 0:
                return torch.round(x)
            else:
                return torch.round(x, decimals=self.decimals)

    @tvm.script.ir_module
    class Expected1:
        @R.function
        def main(
            inp_0: R.Tensor((3, 4), dtype="float32"),
        ) -> R.Tensor((3, 4), dtype="float32"):
            with R.dataflow():
                lv: R.Tensor((3, 4), dtype="float32") = R.round(inp_0)
                gv: R.Tensor((3, 4), dtype="float32") = lv
                R.output(gv)
            return gv

    @tvm.script.ir_module
    class Expected2:
        @R.function
        def main(
            inp_0: R.Tensor((3, 4), dtype="float32"),
        ) -> R.Tensor((3, 4), dtype="float32"):
            with R.dataflow():
                lv: R.Tensor((3, 4), dtype="float32") = R.multiply(inp_0, R.const(100.0, "float32"))
                lv1: R.Tensor((3, 4), dtype="float32") = R.round(lv)
                lv2: R.Tensor((3, 4), dtype="float32") = R.divide(lv1, R.const(100.0, "float32"))
                gv: R.Tensor((3, 4), dtype="float32") = lv2
                R.output(gv)
            return gv

    rounds = [
        (0, Expected1),
        (2, Expected2),
    ]

    for decimals, expected in rounds:
        verify_model(Round(decimals), input_info, {}, expected)

    # Test numerical accuracy with decimals
    test_data = torch.tensor(
        [
            [1.2345, 2.3456, 3.4567, 4.5678],
            [5.6789, 6.7890, 7.8901, 8.9012],
            [9.1234, 10.2345, 11.3456, 12.4567],
        ]
    )

    for decimals in [0, 1, 2, 3]:
        torch_model = Round(decimals)
        graph_model = fx.symbolic_trace(torch_model)
        with torch.no_grad():
            mod = from_fx(graph_model, input_info)

        target = tvm.target.Target("llvm")
        ex = relax.build(mod, target)
        vm = relax.VirtualMachine(ex, tvm.cpu())

        torch_result = torch_model(test_data).numpy()
        tvm_input = tvm.runtime.tensor(test_data.numpy())
        tvm_result = vm["main"](tvm_input).numpy()

        # Use relaxed tolerance due to floating-point precision in decimal operations
        tvm.testing.assert_allclose(tvm_result, torch_result, rtol=1e-3, atol=1e-3)


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