# 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: E501, F841


import tvm
import tvm.testing
from tvm.relax.transform import LegalizeOps
from tvm.script import ir as I
from tvm.script import relax as R
from tvm.script import tirx as T

##################### Search #####################


def test_where():
    # fmt: off
    @tvm.script.ir_module
    class Where:
        @R.function
        def main(condition: R.Tensor((3, 2, 1), "bool"), x: R.Tensor((2, 3), "float32"), y: R.Tensor((2, 1), "float32")) -> R.Tensor((3, 2, 3), "float32"):
            gv: R.Tensor((3, 2, 3), "float32") = R.where(condition, x, y)
            return gv

    @tvm.script.ir_module
    class Expected:
        @R.function
        def main(condition: R.Tensor((3, 2, 1), "bool"), x: R.Tensor((2, 3), "float32"), y: R.Tensor((2, 1), "float32")) -> R.Tensor((3, 2, 3), "float32"):
            gv = R.call_tir(Expected.where, (condition, x, y), R.Tensor((3, 2, 3), dtype="float32"))
            return gv

        @T.prim_func(private=True, s_tir=True)
        def where(rxplaceholder: T.Buffer((T.int64(3), T.int64(2), T.int64(1)), "bool"), rxplaceholder_1: T.Buffer((T.int64(2), T.int64(3)), "float32"), rxplaceholder_2: T.Buffer((T.int64(2), T.int64(1)), "float32"), T_where: T.Buffer((T.int64(3), T.int64(2), T.int64(3)), "float32")):
            T.func_attr({"tirx.noalias": True})
            for i0, i1, i2 in T.grid(T.int64(3), T.int64(2), T.int64(3)):
                with T.sblock("T_where"):
                    ax0, ax1, ax2 = T.axis.remap("SSS", [i0, i1, i2])
                    T.reads(rxplaceholder[ax0, ax1, T.int64(0)], rxplaceholder_1[ax1, ax2], rxplaceholder_2[ax1, T.int64(0)])
                    T.writes(T_where[ax0, ax1, ax2])
                    T_where[ax0, ax1, ax2] = T.Select(T.int64(0) < T.Cast("int64", rxplaceholder[ax0, ax1, T.int64(0)]), rxplaceholder_1[ax1, ax2], rxplaceholder_2[ax1, T.int64(0)])
    # fmt: on

    mod = LegalizeOps()(Where)
    tvm.ir.assert_structural_equal(mod, Expected)


def test_where_symbolic():
    # fmt: off
    @tvm.script.ir_module
    class Where:
        @R.function
        def main(condition: R.Tensor(("a", "b", 1), "bool"), x: R.Tensor(("b", "c"), "float32"), y: R.Tensor(("b", 1), "float32")) -> R.Tensor(("a", "b", "c"), "float32"):
            a = T.int64()
            b = T.int64()
            c = T.int64()
            gv: R.Tensor((a, b, c), "float32") = R.where(condition, x, y)
            return gv

    @tvm.script.ir_module
    class Expected:
        @R.function
        def main(condition: R.Tensor(("a", "b", 1), "bool"), x: R.Tensor(("b", "c"), "float32"), y: R.Tensor(("b", 1), "float32")) -> R.Tensor(("a", "b", "c"), "float32"):
            a = T.int64()
            b = T.int64()
            c = T.int64()
            gv = R.call_tir(Expected.where, (condition, x, y), R.Tensor((a, b, c), dtype="float32"))
            return gv

        @T.prim_func(private=True, s_tir=True)
        def where(var_rxplaceholder: T.handle, var_rxplaceholder_1: T.handle, var_rxplaceholder_2: T.handle, var_T_where: T.handle):
            T.func_attr({"tirx.noalias": True})
            a = T.int64()
            b = T.int64()
            c = T.int64()
            rxplaceholder = T.match_buffer(var_rxplaceholder, [a, b, T.int64(1)], dtype="bool")
            rxplaceholder_1 = T.match_buffer(var_rxplaceholder_1, [b, c], dtype="float32")
            rxplaceholder_2 = T.match_buffer(var_rxplaceholder_2, [b, T.int64(1)], dtype="float32")
            T_where = T.match_buffer(var_T_where, [a, b, c], dtype="float32")
            for i0, i1, i2 in T.grid(a, b, c):
                with T.sblock("T_where"):
                    ax0, ax1, ax2 = T.axis.remap("SSS", [i0, i1, i2])
                    T.reads(rxplaceholder[ax0, ax1, T.int64(0)], rxplaceholder_1[ax1, ax2], rxplaceholder_2[ax1, T.int64(0)])
                    T.writes(T_where[ax0, ax1, ax2])
                    T_where[ax0, ax1, ax2] = T.Select(T.int64(0) < T.Cast("int64", rxplaceholder[ax0, ax1, T.int64(0)]), rxplaceholder_1[ax1, ax2], rxplaceholder_2[ax1, T.int64(0)])
    # fmt: on

    mod = LegalizeOps()(Where)
    tvm.ir.assert_structural_equal(mod, Expected)


def test_argmax():
    # fmt: off
    @tvm.script.ir_module
    class Argmax:
        @R.function
        def main(x: R.Tensor((2, 3, 4, 5), "float32")) -> R.Tensor((2, 4, 5), "int64"):
            gv: R.Tensor((2, 4, 5), "int64") = R.argmax(x, axis=1)
            return gv

    @tvm.script.ir_module
    class Expected:
        @R.function
        def main(x: R.Tensor((2, 3, 4, 5), dtype="float32")) -> R.Tensor((2, 4, 5), dtype="int64"):
            gv = R.call_tir(Expected.argmax, (x,), out_ty=R.Tensor((2, 4, 5), dtype="int64"))
            return gv

        @T.prim_func(private=True, s_tir=True)
        def argmax(rxplaceholder: T.Buffer((T.int64(2), T.int64(3), T.int64(4), T.int64(5)), "float32"), rxplaceholder_red: T.Buffer((T.int64(2), T.int64(4), T.int64(5)), "int64")):
            T.func_attr({"tirx.noalias": True})
            rxplaceholder_red_temp_v0 = T.sblock_alloc_buffer((T.int64(2), T.int64(4), T.int64(5)), "int64")
            rxplaceholder_red_temp_v1 = T.sblock_alloc_buffer((T.int64(2), T.int64(4), T.int64(5)))
            for ax0, ax1, ax2, k1 in T.grid(T.int64(2), T.int64(4), T.int64(5), T.int64(3)):
                with T.sblock("rxplaceholder_red_temp"):
                    v_ax0, v_ax1, v_ax2, v_k1 = T.axis.remap("SSSR", [ax0, ax1, ax2, k1])
                    T.reads(rxplaceholder[v_ax0, v_k1, v_ax1, v_ax2])
                    T.writes(rxplaceholder_red_temp_v0[v_ax0, v_ax1, v_ax2], rxplaceholder_red_temp_v1[v_ax0, v_ax1, v_ax2])
                    with T.init():
                        rxplaceholder_red_temp_v0[v_ax0, v_ax1, v_ax2] = T.int64(-1)
                        rxplaceholder_red_temp_v1[v_ax0, v_ax1, v_ax2] = T.min_value("float32")
                    v_rxplaceholder_red_temp_v0: T.let[T.int64] = T.Select(rxplaceholder_red_temp_v1[v_ax0, v_ax1, v_ax2] > rxplaceholder[v_ax0, v_k1, v_ax1, v_ax2] or (rxplaceholder_red_temp_v1[v_ax0, v_ax1, v_ax2] == rxplaceholder[v_ax0, v_k1, v_ax1, v_ax2] and rxplaceholder_red_temp_v0[v_ax0, v_ax1, v_ax2] < v_k1), rxplaceholder_red_temp_v0[v_ax0, v_ax1, v_ax2], v_k1)
                    v_rxplaceholder_red_temp_v1: T.let[T.float32] = T.Select(rxplaceholder_red_temp_v1[v_ax0, v_ax1, v_ax2] > rxplaceholder[v_ax0, v_k1, v_ax1, v_ax2], rxplaceholder_red_temp_v1[v_ax0, v_ax1, v_ax2], rxplaceholder[v_ax0, v_k1, v_ax1, v_ax2])
                    rxplaceholder_red_temp_v0[v_ax0, v_ax1, v_ax2] = v_rxplaceholder_red_temp_v0
                    rxplaceholder_red_temp_v1[v_ax0, v_ax1, v_ax2] = v_rxplaceholder_red_temp_v1
            for ax0, ax1, ax2 in T.grid(T.int64(2), T.int64(4), T.int64(5)):
                with T.sblock("rxplaceholder_red"):
                    v_ax0, v_ax1, v_ax2 = T.axis.remap("SSS", [ax0, ax1, ax2])
                    T.reads(rxplaceholder_red_temp_v0[v_ax0, v_ax1, v_ax2])
                    T.writes(rxplaceholder_red[v_ax0, v_ax1, v_ax2])
                    rxplaceholder_red[v_ax0, v_ax1, v_ax2] = rxplaceholder_red_temp_v0[v_ax0, v_ax1, v_ax2]
    # fmt: on

    mod = LegalizeOps()(Argmax)
    tvm.ir.assert_structural_equal(mod, Expected)


def test_argmax_symbolic():
    # fmt: off
    @tvm.script.ir_module
    class Argmax:
        @R.function
        def main(x: R.Tensor(("a", "b", "c", "d"), "float32")) -> R.Tensor(("a", 1, "c", "d"), "int64"):
            a = T.int64()
            c = T.int64()
            d = T.int64()
            gv: R.Tensor((a, 1, c, d), "int64") = R.argmax(x, axis=1, keepdims=True)
            return gv

    @tvm.script.ir_module
    class Expected:
        @R.function
        def main(x: R.Tensor(("a", "b", "c", "d"), dtype="float32")) -> R.Tensor(("a", 1, "c", "d"), dtype="int64"):
            a = T.int64()
            c = T.int64()
            d = T.int64()
            gv = R.call_tir(Expected.argmax, (x,), out_ty=R.Tensor((a, 1, c, d), dtype="int64"))
            return gv

        @T.prim_func(private=True, s_tir=True)
        def argmax(var_rxplaceholder: T.handle, var_rxplaceholder_red: T.handle):
            T.func_attr({"tirx.noalias": True})
            a = T.int64()
            b = T.int64()
            c = T.int64()
            d = T.int64()
            rxplaceholder = T.match_buffer(var_rxplaceholder, (a, b, c, d))
            rxplaceholder_red = T.match_buffer(var_rxplaceholder_red, (a, T.int64(1), c, d), "int64")
            # with T.sblock("root"):
            rxplaceholder_red_temp_v0 = T.sblock_alloc_buffer((a, T.int64(1), c, d), "int64")
            rxplaceholder_red_temp_v1 = T.sblock_alloc_buffer((a, T.int64(1), c, d))
            for ax0, ax1, ax2, ax3, k1 in T.grid(a, T.int64(1), c, d, b):
                with T.sblock("rxplaceholder_red_temp"):
                    v_ax0, v_ax1, v_ax2, v_ax3, v_k1 = T.axis.remap("SSSSR", [ax0, ax1, ax2, ax3, k1])
                    T.reads(rxplaceholder[v_ax0, v_k1, v_ax2, v_ax3])
                    T.writes(rxplaceholder_red_temp_v0[v_ax0, v_ax1, v_ax2, v_ax3], rxplaceholder_red_temp_v1[v_ax0, v_ax1, v_ax2, v_ax3])
                    with T.init():
                        rxplaceholder_red_temp_v0[v_ax0, v_ax1, v_ax2, v_ax3] = T.int64(-1)
                        rxplaceholder_red_temp_v1[v_ax0, v_ax1, v_ax2, v_ax3] = T.min_value("float32")
                    v_rxplaceholder_red_temp_v0: T.let[T.int64] = T.Select(rxplaceholder_red_temp_v1[v_ax0, v_ax1, v_ax2, v_ax3] > rxplaceholder[v_ax0, v_k1, v_ax2, v_ax3] or (rxplaceholder_red_temp_v1[v_ax0, v_ax1, v_ax2, v_ax3] == rxplaceholder[v_ax0, v_k1, v_ax2, v_ax3] and rxplaceholder_red_temp_v0[v_ax0, v_ax1, v_ax2, v_ax3] < v_k1), rxplaceholder_red_temp_v0[v_ax0, v_ax1, v_ax2, v_ax3], v_k1)
                    v_rxplaceholder_red_temp_v1: T.let[T.float32] = T.Select(rxplaceholder_red_temp_v1[v_ax0, v_ax1, v_ax2, v_ax3] > rxplaceholder[v_ax0, v_k1, v_ax2, v_ax3], rxplaceholder_red_temp_v1[v_ax0, v_ax1, v_ax2, v_ax3], rxplaceholder[v_ax0, v_k1, v_ax2, v_ax3])
                    rxplaceholder_red_temp_v0[v_ax0, v_ax1, v_ax2, v_ax3] = v_rxplaceholder_red_temp_v0
                    rxplaceholder_red_temp_v1[v_ax0, v_ax1, v_ax2, v_ax3] = v_rxplaceholder_red_temp_v1
            for ax0, ax1, ax2, ax3 in T.grid(a, T.int64(1), c, d):
                with T.sblock("rxplaceholder_red"):
                    v_ax0, v_ax1, v_ax2, v_ax3 = T.axis.remap("SSSS", [ax0, ax1, ax2, ax3])
                    T.reads(rxplaceholder_red_temp_v0[v_ax0, v_ax1, v_ax2, v_ax3])
                    T.writes(rxplaceholder_red[v_ax0, v_ax1, v_ax2, v_ax3])
                    rxplaceholder_red[v_ax0, v_ax1, v_ax2, v_ax3] = rxplaceholder_red_temp_v0[v_ax0, v_ax1, v_ax2, v_ax3]
    # fmt: on

    mod = LegalizeOps()(Argmax)
    tvm.ir.assert_structural_equal(mod, Expected)


def test_argmin():
    # fmt: off
    @tvm.script.ir_module
    class Argmin:
        @R.function
        def main(x: R.Tensor((2, 3, 4, 5), "float32")) -> R.Tensor((), "int64"):
            gv: R.Tensor((), "int64") = R.argmin(x)
            return gv

    @tvm.script.ir_module
    class Expected:
        @T.prim_func(private=True, s_tir=True)
        def argmin(rxplaceholder: T.Buffer((T.int64(2), T.int64(3), T.int64(4), T.int64(5)), "float32"), rxplaceholder_red: T.Buffer((), "int64")):
            T.func_attr({"tirx.noalias": True})
            rxplaceholder_red_temp_v0 = T.sblock_alloc_buffer((), "int64")
            rxplaceholder_red_temp_v1 = T.sblock_alloc_buffer(())
            for k0, k1, k2, k3 in T.grid(T.int64(2), T.int64(3), T.int64(4), T.int64(5)):
                with T.sblock("rxplaceholder_red_temp"):
                    v_k0, v_k1, v_k2, v_k3 = T.axis.remap("RRRR", [k0, k1, k2, k3])
                    T.reads(rxplaceholder[v_k0, v_k1, v_k2, v_k3])
                    T.writes(rxplaceholder_red_temp_v0[()], rxplaceholder_red_temp_v1[()])
                    with T.init():
                        rxplaceholder_red_temp_v0[()] = T.int64(-1)
                        rxplaceholder_red_temp_v1[()] = T.max_value("float32")
                    v_rxplaceholder_red_temp_v0: T.let[T.int64] = T.Select(rxplaceholder_red_temp_v1[()] < rxplaceholder[v_k0, v_k1, v_k2, v_k3] or (rxplaceholder_red_temp_v1[()] == rxplaceholder[v_k0, v_k1, v_k2, v_k3] and rxplaceholder_red_temp_v0[()] < v_k0 * T.int64(60) + v_k1 * T.int64(20) + v_k2 * T.int64(5) + v_k3), rxplaceholder_red_temp_v0[()], v_k0 * T.int64(60) + v_k1 * T.int64(20) + v_k2 * T.int64(5) + v_k3)
                    v_rxplaceholder_red_temp_v1: T.let[T.float32] = T.Select(rxplaceholder_red_temp_v1[()] < rxplaceholder[v_k0, v_k1, v_k2, v_k3], rxplaceholder_red_temp_v1[()], rxplaceholder[v_k0, v_k1, v_k2, v_k3])
                    rxplaceholder_red_temp_v0[()] = v_rxplaceholder_red_temp_v0
                    rxplaceholder_red_temp_v1[()] = v_rxplaceholder_red_temp_v1
            with T.sblock("rxplaceholder_red"):
                vi = T.axis.spatial(1, T.int64(0))
                T.reads(rxplaceholder_red_temp_v0[()])
                T.writes(rxplaceholder_red[()])
                rxplaceholder_red[()] = rxplaceholder_red_temp_v0[()]

        @R.function
        def main(x: R.Tensor((2, 3, 4, 5), dtype="float32")) -> R.Tensor((), dtype="int64"):
            gv = R.call_tir(Expected.argmin, (x,), out_ty=R.Tensor((), dtype="int64"))
            return gv
    # fmt: on

    mod = LegalizeOps()(Argmin)
    tvm.ir.assert_structural_equal(mod, Expected)


def test_argmin_symbolic():
    # fmt: off
    @tvm.script.ir_module
    class Argmin:
        @R.function
        def main(x: R.Tensor(("a", "b", "c", "d"), "float32")) -> R.Tensor((1, 1, 1, 1), "int64"):
            gv: R.Tensor((1, 1, 1, 1), "int64") = R.argmin(x, keepdims=True)
            return gv

    @tvm.script.ir_module
    class Expected:
        @T.prim_func(private=True, s_tir=True)
        def argmin(var_rxplaceholder: T.handle, rxplaceholder_red: T.Buffer((T.int64(1), T.int64(1), T.int64(1), T.int64(1)), "int64")):
            T.func_attr({"tirx.noalias": True})
            a = T.int64()
            b = T.int64()
            c = T.int64()
            d = T.int64()
            rxplaceholder = T.match_buffer(var_rxplaceholder, (a, b, c, d))
            rxplaceholder_red_temp_v0 = T.sblock_alloc_buffer((T.int64(1), T.int64(1), T.int64(1), T.int64(1)), "int64")
            rxplaceholder_red_temp_v1 = T.sblock_alloc_buffer((T.int64(1), T.int64(1), T.int64(1), T.int64(1)))
            for ax0, ax1, ax2, ax3, k0, k1, k2, k3 in T.grid(T.int64(1), T.int64(1), T.int64(1), T.int64(1), a, b, c, d):
                with T.sblock("rxplaceholder_red_temp"):
                    v_ax0, v_ax1, v_ax2, v_ax3, v_k0, v_k1, v_k2, v_k3 = T.axis.remap("SSSSRRRR", [ax0, ax1, ax2, ax3, k0, k1, k2, k3])
                    T.reads(rxplaceholder[v_k0, v_k1, v_k2, v_k3])
                    T.writes(rxplaceholder_red_temp_v0[v_ax0, v_ax1, v_ax2, v_ax3], rxplaceholder_red_temp_v1[v_ax0, v_ax1, v_ax2, v_ax3])
                    with T.init():
                        rxplaceholder_red_temp_v0[v_ax0, v_ax1, v_ax2, v_ax3] = T.int64(-1)
                        rxplaceholder_red_temp_v1[v_ax0, v_ax1, v_ax2, v_ax3] = T.max_value("float32")
                    v_rxplaceholder_red_temp_v0: T.let[T.int64] = T.Select(rxplaceholder_red_temp_v1[v_ax0, v_ax1, v_ax2, v_ax3] < rxplaceholder[v_k0, v_k1, v_k2, v_k3] or (rxplaceholder_red_temp_v1[v_ax0, v_ax1, v_ax2, v_ax3] == rxplaceholder[v_k0, v_k1, v_k2, v_k3] and rxplaceholder_red_temp_v0[v_ax0, v_ax1, v_ax2, v_ax3] < ((v_k0 * b + v_k1) * c + v_k2) * d + v_k3), rxplaceholder_red_temp_v0[v_ax0, v_ax1, v_ax2, v_ax3], ((v_k0 * b + v_k1) * c + v_k2) * d + v_k3)
                    v_rxplaceholder_red_temp_v1: T.let[T.float32] = T.Select(rxplaceholder_red_temp_v1[v_ax0, v_ax1, v_ax2, v_ax3] < rxplaceholder[v_k0, v_k1, v_k2, v_k3], rxplaceholder_red_temp_v1[v_ax0, v_ax1, v_ax2, v_ax3], rxplaceholder[v_k0, v_k1, v_k2, v_k3])
                    rxplaceholder_red_temp_v0[v_ax0, v_ax1, v_ax2, v_ax3] = v_rxplaceholder_red_temp_v0
                    rxplaceholder_red_temp_v1[v_ax0, v_ax1, v_ax2, v_ax3] = v_rxplaceholder_red_temp_v1
            for ax0, ax1, ax2, ax3 in T.grid(T.int64(1), T.int64(1), T.int64(1), T.int64(1)):
                with T.sblock("rxplaceholder_red"):
                    v_ax0, v_ax1, v_ax2, v_ax3 = T.axis.remap("SSSS", [ax0, ax1, ax2, ax3])
                    T.reads(rxplaceholder_red_temp_v0[v_ax0, v_ax1, v_ax2, v_ax3])
                    T.writes(rxplaceholder_red[v_ax0, v_ax1, v_ax2, v_ax3])
                    rxplaceholder_red[v_ax0, v_ax1, v_ax2, v_ax3] = rxplaceholder_red_temp_v0[v_ax0, v_ax1, v_ax2, v_ax3]

        @R.function
        def main(x: R.Tensor(("a", "b", "c", "d"), dtype="float32")) -> R.Tensor((1, 1, 1, 1), dtype="int64"):
            gv = R.call_tir(Expected.argmin, (x,), out_ty=R.Tensor((1, 1, 1, 1), dtype="int64"))
            return gv
    # fmt: on

    mod = LegalizeOps()(Argmin)
    tvm.ir.assert_structural_equal(mod, Expected)


##################### Statistical #####################


def test_max():
    # fmt: off
    @tvm.script.ir_module
    class Max:
        @R.function
        def main(x: R.Tensor((2, 3, 4, 5), "float32")) -> R.Tensor((2, 5), "float32"):
            gv: R.Tensor((2, 5), "float32") = R.max(x, axis=[1, 2])
            return gv

    @tvm.script.ir_module
    class Expected:
        @R.function
        def main(x: R.Tensor((2, 3, 4, 5), "float32")) -> R.Tensor((2, 5), "float32"):
            gv = R.call_tir(Expected.max, (x,), R.Tensor((2, 5), dtype="float32"))
            return gv

        @T.prim_func(private=True, s_tir=True)
        def max(rxplaceholder: T.Buffer((T.int64(2), T.int64(3), T.int64(4), T.int64(5)), "float32"), rxplaceholder_red: T.Buffer((T.int64(2), T.int64(5)), "float32")):
            T.func_attr({"tirx.noalias": True})
            for i0, i1, i2, i3 in T.grid(T.int64(2), T.int64(5), T.int64(3), T.int64(4)):
                with T.sblock("rxplaceholder_red"):
                    ax0, ax1, k1, k2 = T.axis.remap("SSRR", [i0, i1, i2, i3])
                    T.reads(rxplaceholder[ax0, k1, k2, ax1])
                    T.writes(rxplaceholder_red[ax0, ax1])
                    with T.init():
                        rxplaceholder_red[ax0, ax1] = T.min_value("float32")
                    rxplaceholder_red[ax0, ax1] = T.max(rxplaceholder_red[ax0, ax1], rxplaceholder[ax0, k1, k2, ax1])
    # fmt: on

    mod = LegalizeOps()(Max)
    tvm.ir.assert_structural_equal(mod, Expected)


def test_max_symbolic():
    # fmt: off
    @tvm.script.ir_module
    class Max:
        @R.function
        def main(x: R.Tensor(("a", "b", "c", "d"), "float32")) -> R.Tensor(("a", "d"), "float32"):
            a = T.int64()
            d = T.int64()
            gv: R.Tensor((a, d), "float32") = R.max(x, axis=[1, 2])
            return gv

    @tvm.script.ir_module
    class Expected:
        @R.function
        def main(x: R.Tensor(("a", "b", "c", "d"), "float32")) -> R.Tensor(("a", "d"), "float32"):
            a = T.int64()
            d = T.int64()
            gv = R.call_tir(Expected.max, (x,), R.Tensor((a, d), dtype="float32"))
            return gv

        @T.prim_func(private=True, s_tir=True)
        def max(var_rxplaceholder: T.handle, var_rxplaceholder_red: T.handle):
            T.func_attr({"tirx.noalias": True})
            a = T.int64()
            b = T.int64()
            c = T.int64()
            d = T.int64()
            rxplaceholder = T.match_buffer(var_rxplaceholder, [a, b, c, d], dtype="float32")
            rxplaceholder_red = T.match_buffer(var_rxplaceholder_red, [a, d], dtype="float32")
            for i0, i1, i2, i3 in T.grid(a, d, b, c):
                with T.sblock("rxplaceholder_red"):
                    ax0, ax1, k1, k2 = T.axis.remap("SSRR", [i0, i1, i2, i3])
                    T.reads(rxplaceholder[ax0, k1, k2, ax1])
                    T.writes(rxplaceholder_red[ax0, ax1])
                    with T.init():
                        rxplaceholder_red[ax0, ax1] = T.min_value("float32")
                    rxplaceholder_red[ax0, ax1] = T.max(rxplaceholder_red[ax0, ax1], rxplaceholder[ax0, k1, k2, ax1])
    # fmt: on

    mod = LegalizeOps()(Max)
    tvm.ir.assert_structural_equal(mod, Expected)


def test_min():
    # fmt: off
    @tvm.script.ir_module
    class Min:
        @R.function
        def main(x: R.Tensor((2, 3, 4, 5), "float32")) -> R.Tensor((2, 1, 1, 5), "float32"):
            gv: R.Tensor((2, 1, 1, 5), "float32") = R.min(x, axis=[1, 2], keepdims=True)
            return gv

    @tvm.script.ir_module
    class Expected:
        @R.function
        def main(x: R.Tensor((2, 3, 4, 5), "float32")) -> R.Tensor((2, 1, 1, 5), "float32"):
            gv = R.call_tir(Expected.min, (x,), R.Tensor((2, 1, 1, 5), dtype="float32"))
            return gv

        @T.prim_func(private=True, s_tir=True)
        def min(rxplaceholder: T.Buffer((T.int64(2), T.int64(3), T.int64(4), T.int64(5)), "float32"), rxplaceholder_red: T.Buffer((T.int64(2), T.int64(1), T.int64(1), T.int64(5)), "float32")):
            T.func_attr({"tirx.noalias": True})
            for i0, i1, i2, i3, i4, i5 in T.grid(T.int64(2), T.int64(1), T.int64(1), T.int64(5), T.int64(3), T.int64(4)):
                with T.sblock("rxplaceholder_red"):
                    ax0, ax1, ax2, ax3, k1, k2 = T.axis.remap("SSSSRR", [i0, i1, i2, i3, i4, i5])
                    T.reads(rxplaceholder[ax0, k1, k2, ax3])
                    T.writes(rxplaceholder_red[ax0, ax1, ax2, ax3])
                    with T.init():
                        rxplaceholder_red[ax0, ax1, ax2, ax3] = T.max_value("float32")
                    rxplaceholder_red[ax0, ax1, ax2, ax3] = T.min(rxplaceholder_red[ax0, ax1, ax2, ax3], rxplaceholder[ax0, k1, k2, ax3])
    # fmt: on

    mod = LegalizeOps()(Min)
    tvm.ir.assert_structural_equal(mod, Expected)


def test_min_symbolic():
    # fmt: off
    @tvm.script.ir_module
    class Min:
        @R.function
        def main(x: R.Tensor(("a", "b", "c", "d"), "float32")) -> R.Tensor(("a", 1, 1, "d"), "float32"):
            a = T.int64()
            d = T.int64()
            gv: R.Tensor((a, 1, 1, d), "float32") = R.min(x, axis=[1, 2], keepdims=True)
            return gv

    @tvm.script.ir_module
    class Expected:
        @R.function
        def main(x: R.Tensor(("a", "b", "c", "d"), "float32")) -> R.Tensor(("a", 1, 1, "d"), "float32"):
            a = T.int64()
            d = T.int64()
            gv = R.call_tir(Expected.min, (x,), R.Tensor((a, 1, 1, d), dtype="float32"))
            return gv

        @T.prim_func(private=True, s_tir=True)
        def min(var_rxplaceholder: T.handle, var_rxplaceholder_red: T.handle):
            T.func_attr({"tirx.noalias": True})
            a = T.int64()
            b = T.int64()
            c = T.int64()
            d = T.int64()
            rxplaceholder = T.match_buffer(var_rxplaceholder, [a, b, c, d], dtype="float32")
            rxplaceholder_red = T.match_buffer(var_rxplaceholder_red, [a, T.int64(1), T.int64(1), d], dtype="float32")
            for i0, i1, i2, i3, i4, i5 in T.grid(a, T.int64(1), T.int64(1), d, b, c):
                with T.sblock("rxplaceholder_red"):
                    ax0, ax1, ax2, ax3, k1, k2 = T.axis.remap("SSSSRR", [i0, i1, i2, i3, i4, i5])
                    T.reads(rxplaceholder[ax0, k1, k2, ax3])
                    T.writes(rxplaceholder_red[ax0, ax1, ax2, ax3])
                    with T.init():
                        rxplaceholder_red[ax0, ax1, ax2, ax3] = T.max_value("float32")
                    rxplaceholder_red[ax0, ax1, ax2, ax3] = T.min(rxplaceholder_red[ax0, ax1, ax2, ax3], rxplaceholder[ax0, k1, k2, ax3])
    # fmt: on

    mod = LegalizeOps()(Min)
    tvm.ir.assert_structural_equal(mod, Expected)


def test_sum():
    # fmt: off
    @tvm.script.ir_module
    class Sum:
        @R.function
        def main(x: R.Tensor((2, 3, 4, 5), "float32")) -> R.Tensor((), "float32"):
            gv: R.Tensor((), "float32") = R.sum(x)
            return gv

    @tvm.script.ir_module
    class Expected:
        @R.function
        def main(x: R.Tensor((2, 3, 4, 5), "float32")) -> R.Tensor((), "float32"):
            gv = R.call_tir(Expected.sum, (x,), R.Tensor((), dtype="float32"))
            return gv

        @T.prim_func(private=True, s_tir=True)
        def sum(rxplaceholder: T.Buffer((T.int64(2), T.int64(3), T.int64(4), T.int64(5)), "float32"), rxplaceholder_red: T.Buffer((), "float32")):
            T.func_attr({"tirx.noalias": True})
            for i0, i1, i2, i3 in T.grid(T.int64(2), T.int64(3), T.int64(4), T.int64(5)):
                with T.sblock("rxplaceholder_red"):
                    k0, k1, k2, k3 = T.axis.remap("RRRR", [i0, i1, i2, i3])
                    T.reads(rxplaceholder[k0, k1, k2, k3])
                    T.writes(rxplaceholder_red[()])
                    with T.init():
                        rxplaceholder_red[()] = T.float32(0)
                    rxplaceholder_red[()] = rxplaceholder_red[()] + rxplaceholder[k0, k1, k2, k3]
    # fmt: on

    mod = LegalizeOps()(Sum)
    tvm.ir.assert_structural_equal(mod, Expected)


def test_sum_symbolic():
    # fmt: off
    @tvm.script.ir_module
    class Sum:
        @R.function
        def main(x: R.Tensor(("a", "b", "c", "d"), "float32")) -> R.Tensor((), "float32"):
            gv: R.Tensor((), "float32") = R.sum(x)
            return gv

    @tvm.script.ir_module
    class Expected:
        @R.function
        def main(x: R.Tensor(("a", "b", "c", "d"), "float32")) -> R.Tensor((), "float32"):
            gv = R.call_tir(Expected.sum, (x,), R.Tensor((), dtype="float32"))
            return gv

        @T.prim_func(private=True, s_tir=True)
        def sum(var_rxplaceholder: T.handle, rxplaceholder_red: T.Buffer((), "float32")):
            T.func_attr({"tirx.noalias": True})
            a = T.int64()
            b = T.int64()
            c = T.int64()
            d = T.int64()
            rxplaceholder = T.match_buffer(var_rxplaceholder, [a, b, c, d], dtype="float32")
            for i0, i1, i2, i3 in T.grid(a, b, c, d):
                with T.sblock("rxplaceholder_red"):
                    k0, k1, k2, k3 = T.axis.remap("RRRR", [i0, i1, i2, i3])
                    T.reads(rxplaceholder[k0, k1, k2, k3])
                    T.writes(rxplaceholder_red[()])
                    with T.init():
                        rxplaceholder_red[()] = T.float32(0)
                    rxplaceholder_red[()] = rxplaceholder_red[()] + rxplaceholder[k0, k1, k2, k3]
    # fmt: on

    mod = LegalizeOps()(Sum)
    tvm.ir.assert_structural_equal(mod, Expected)


def test_prod():
    # fmt: off
    @tvm.script.ir_module
    class Prod:
        @R.function
        def main(x: R.Tensor((2, 3, 4, 5), "float32")) -> R.Tensor((1, 1, 1, 1), "float32"):
            gv: R.Tensor((1, 1, 1, 1), "float32") = R.prod(x, keepdims=True)
            return gv

    @tvm.script.ir_module
    class Expected:
        @R.function
        def main(x: R.Tensor((2, 3, 4, 5), "float32")) -> R.Tensor((1, 1, 1, 1), "float32"):
            gv = R.call_tir(Expected.prod, (x,), R.Tensor((1, 1, 1, 1), dtype="float32"))
            return gv

        @T.prim_func(private=True, s_tir=True)
        def prod(rxplaceholder: T.Buffer((T.int64(2), T.int64(3), T.int64(4), T.int64(5)), "float32"), rxplaceholder_red: T.Buffer((T.int64(1), T.int64(1), T.int64(1), T.int64(1)), "float32")):
            T.func_attr({"tirx.noalias": True})
            for i0, i1, i2, i3, i4, i5, i6, i7 in T.grid(T.int64(1), T.int64(1), T.int64(1), T.int64(1), T.int64(2), T.int64(3), T.int64(4), T.int64(5)):
                with T.sblock("rxplaceholder_red"):
                    ax0, ax1, ax2, ax3, k0, k1, k2, k3 = T.axis.remap("SSSSRRRR", [i0, i1, i2, i3, i4, i5, i6, i7])
                    T.reads(rxplaceholder[k0, k1, k2, k3])
                    T.writes(rxplaceholder_red[ax0, ax1, ax2, ax3])
                    with T.init():
                        rxplaceholder_red[ax0, ax1, ax2, ax3] = T.float32(1)
                    rxplaceholder_red[ax0, ax1, ax2, ax3] = rxplaceholder_red[ax0, ax1, ax2, ax3] * rxplaceholder[k0, k1, k2, k3]
    # fmt: on

    mod = LegalizeOps()(Prod)
    tvm.ir.assert_structural_equal(mod, Expected)


def test_prod_bool():
    # fmt: off
    @tvm.script.ir_module
    class Prod:
        @R.function
        def main(x: R.Tensor((2, 3, 4, 5), "bool")) -> R.Tensor((1, 1, 1, 1), "bool"):
            gv: R.Tensor((1, 1, 1, 1), "bool") = R.prod(x, keepdims=True)
            return gv

    @tvm.script.ir_module
    class Expected:
        @R.function
        def main(x: R.Tensor((2, 3, 4, 5), "bool")) -> R.Tensor((1, 1, 1, 1), "bool"):
            gv = R.call_tir(Expected.prod, (x,), R.Tensor((1, 1, 1, 1), dtype="bool"))
            return gv

        @T.prim_func(private=True, s_tir=True)
        def prod(rxplaceholder: T.Buffer((T.int64(2), T.int64(3), T.int64(4), T.int64(5)), "bool"), rxplaceholder_red: T.Buffer((T.int64(1), T.int64(1), T.int64(1), T.int64(1)), "bool")):
            T.func_attr({"tirx.noalias": True})
            for i0, i1, i2, i3, i4, i5, i6, i7 in T.grid(T.int64(1), T.int64(1), T.int64(1), T.int64(1), T.int64(2), T.int64(3), T.int64(4), T.int64(5)):
                with T.sblock("rxplaceholder_red"):
                    ax0, ax1, ax2, ax3, k0, k1, k2, k3 = T.axis.remap("SSSSRRRR", [i0, i1, i2, i3, i4, i5, i6, i7])
                    T.reads(rxplaceholder[k0, k1, k2, k3])
                    T.writes(rxplaceholder_red[ax0, ax1, ax2, ax3])
                    with T.init():
                        rxplaceholder_red[ax0, ax1, ax2, ax3] = T.bool(1)
                    rxplaceholder_red[ax0, ax1, ax2, ax3] = rxplaceholder_red[ax0, ax1, ax2, ax3] and rxplaceholder[k0, k1, k2, k3]
    # fmt: on

    mod = LegalizeOps()(Prod)
    tvm.ir.assert_structural_equal(mod, Expected)


def test_prod_symbolic():
    # fmt: off
    @tvm.script.ir_module
    class Prod:
        @R.function
        def main(x: R.Tensor(("a", "b", "c", "d"), "float32")) -> R.Tensor((1, 1, 1, 1), "float32"):
            gv: R.Tensor((1, 1, 1, 1), "float32") = R.prod(x, keepdims=True)
            return gv

    @tvm.script.ir_module
    class Expected:
        @R.function
        def main(x: R.Tensor(("a", "b", "c", "d"), "float32")) -> R.Tensor((1, 1, 1, 1), "float32"):
            gv = R.call_tir(Expected.prod, (x,), R.Tensor((1, 1, 1, 1), dtype="float32"))
            return gv

        @T.prim_func(private=True, s_tir=True)
        def prod(var_rxplaceholder: T.handle, rxplaceholder_red: T.Buffer((T.int64(1), T.int64(1), T.int64(1), T.int64(1)), "float32")):
            T.func_attr({"tirx.noalias": True})
            a = T.int64()
            b = T.int64()
            c = T.int64()
            d = T.int64()
            rxplaceholder = T.match_buffer(var_rxplaceholder, [a, b, c, d], dtype="float32")
            for i0, i1, i2, i3, i4, i5, i6, i7 in T.grid(T.int64(1), T.int64(1), T.int64(1), T.int64(1), a, b, c, d):
                with T.sblock("rxplaceholder_red"):
                    ax0, ax1, ax2, ax3, k0, k1, k2, k3 = T.axis.remap("SSSSRRRR", [i0, i1, i2, i3, i4, i5, i6, i7])
                    T.reads(rxplaceholder[k0, k1, k2, k3])
                    T.writes(rxplaceholder_red[ax0, ax1, ax2, ax3])
                    with T.init():
                        rxplaceholder_red[ax0, ax1, ax2, ax3] = T.float32(1)
                    rxplaceholder_red[ax0, ax1, ax2, ax3] = rxplaceholder_red[ax0, ax1, ax2, ax3] * rxplaceholder[k0, k1, k2, k3]
    # fmt: on

    mod = LegalizeOps()(Prod)
    tvm.ir.assert_structural_equal(mod, Expected)


def test_sum_zero_dim_axis_identity():
    # fmt: off
    @tvm.script.ir_module
    class Sum:
        @R.function
        def main(x: R.Tensor((2, 0, 4), "float32")) -> R.Tensor((2, 4), "float32"):
            gv: R.Tensor((2, 4), "float32") = R.sum(x, axis=[1], keepdims=False)
            return gv
    # fmt: on

    mod = LegalizeOps()(Sum)
    script = mod.script()
    assert "T.axis.reduce" not in script
    assert "T.float32(0)" in script or "T.float32(0.0)" in script


def test_sum_zero_dim_negative_axis_identity():
    # fmt: off
    @tvm.script.ir_module
    class Sum:
        @R.function
        def main(x: R.Tensor((2, 3, 0), "float32")) -> R.Tensor((2, 3), "float32"):
            gv: R.Tensor((2, 3), "float32") = R.sum(x, axis=[-1], keepdims=False)
            return gv
    # fmt: on

    mod = LegalizeOps()(Sum)
    script = mod.script()
    assert "T.axis.reduce" not in script
    assert "T.float32(0)" in script or "T.float32(0.0)" in script


def test_prod_zero_dim_axis_identity():
    # fmt: off
    @tvm.script.ir_module
    class Prod:
        @R.function
        def main(x: R.Tensor((2, 0, 4), "float32")) -> R.Tensor((2, 4), "float32"):
            gv: R.Tensor((2, 4), "float32") = R.prod(x, axis=[1], keepdims=False)
            return gv
    # fmt: on

    mod = LegalizeOps()(Prod)
    script = mod.script()
    assert "T.axis.reduce" not in script
    assert "T.float32(1)" in script or "T.float32(1.0)" in script


def test_prod_bool_zero_dim_axis_identity():
    # fmt: off
    @tvm.script.ir_module
    class Prod:
        @R.function
        def main(x: R.Tensor((2, 0, 4), "bool")) -> R.Tensor((2, 4), "bool"):
            gv: R.Tensor((2, 4), "bool") = R.prod(x, axis=[1], keepdims=False)
            return gv
    # fmt: on

    mod = LegalizeOps()(Prod)
    script = mod.script()
    assert "T.axis.reduce" not in script
    assert "T.bool(1)" in script or "T.bool(True)" in script


def test_mean():
    # fmt: off
    @tvm.script.ir_module
    class Mean:
        @R.function
        def main(x: R.Tensor((2, 3, 4, 5), "float32")) -> R.Tensor((3, 4), "float32"):
            gv: R.Tensor((3, 4), "float32") = R.mean(x, [0, 3])
            return gv

    @tvm.script.ir_module
    class Expected:
        @R.function
        def main(x: R.Tensor((2, 3, 4, 5), "float32")) -> R.Tensor((3, 4), "float32"):
            gv = R.call_tir(Expected.mean, (x,), R.Tensor((3, 4), dtype="float32"))
            return gv

        @T.prim_func(private=True, s_tir=True)
        def mean(rxplaceholder: T.Buffer((T.int64(2), T.int64(3), T.int64(4), T.int64(5)), "float32"), T_divide: T.Buffer((T.int64(3), T.int64(4)), "float32")):
            T.func_attr({"tirx.noalias": True})
            rxplaceholder_red = T.sblock_alloc_buffer([T.int64(3), T.int64(4)], dtype="float32")
            for i0, i1, i2, i3 in T.grid(T.int64(3), T.int64(4), T.int64(2), T.int64(5)):
                with T.sblock("rxplaceholder_red"):
                    ax0, ax1, k0, k3 = T.axis.remap("SSRR", [i0, i1, i2, i3])
                    T.reads(rxplaceholder[k0, ax0, ax1, k3])
                    T.writes(rxplaceholder_red[ax0, ax1])
                    with T.init():
                        rxplaceholder_red[ax0, ax1] = T.float32(0)
                    rxplaceholder_red[ax0, ax1] = rxplaceholder_red[ax0, ax1] + rxplaceholder[k0, ax0, ax1, k3]
            for i0, i1 in T.grid(T.int64(3), T.int64(4)):
                with T.sblock("T_divide"):
                    ax0, ax1 = T.axis.remap("SS", [i0, i1])
                    T.reads(rxplaceholder_red[ax0, ax1])
                    T.writes(T_divide[ax0, ax1])
                    T_divide[ax0, ax1] = rxplaceholder_red[ax0, ax1] / T.float32(10)
    # fmt: on

    mod = LegalizeOps()(Mean)
    tvm.ir.assert_structural_equal(mod, Expected)


def test_mean_symbolic():
    # fmt: off
    @tvm.script.ir_module
    class Mean:
        @R.function
        def main(x: R.Tensor(("a", "b", "c", "d"), "float32")) -> R.Tensor(("b", "c"), "float32"):
            b = T.int64()
            c = T.int64()
            gv: R.Tensor((b, c), "float32") = R.mean(x, [0, 3])
            return gv

    @tvm.script.ir_module
    class Expected:
        @R.function
        def main(x: R.Tensor(("a", "b", "c", "d"), dtype="float32")) -> R.Tensor(("b", "c"), dtype="float32"):
            b = T.int64()
            c = T.int64()
            gv = R.call_tir(Expected.mean, (x,), R.Tensor((b, c), dtype="float32"))
            return gv

        @T.prim_func(private=True, s_tir=True)
        def mean(var_rxplaceholder: T.handle, var_T_divide: T.handle):
            T.func_attr({"tirx.noalias": True})
            a = T.int64()
            b = T.int64()
            c = T.int64()
            d = T.int64()
            rxplaceholder = T.match_buffer(var_rxplaceholder, [a, b, c, d], dtype="float32")
            T_divide = T.match_buffer(var_T_divide, [b, c], dtype="float32")
            rxplaceholder_red = T.sblock_alloc_buffer([b, c], dtype="float32")
            for i0, i1, i2, i3 in T.grid(b, c, a, d):
                with T.sblock("rxplaceholder_red"):
                    ax0, ax1, k0, k3 = T.axis.remap("SSRR", [i0, i1, i2, i3])
                    T.reads(rxplaceholder[k0, ax0, ax1, k3])
                    T.writes(rxplaceholder_red[ax0, ax1])
                    with T.init():
                        rxplaceholder_red[ax0, ax1] = T.float32(0)
                    rxplaceholder_red[ax0, ax1] = rxplaceholder_red[ax0, ax1] + rxplaceholder[k0, ax0, ax1, k3]
            for i0, i1 in T.grid(b, c):
                with T.sblock("T_divide"):
                    ax0, ax1 = T.axis.remap("SS", [i0, i1])
                    T.reads(rxplaceholder_red[ax0, ax1])
                    T.writes(T_divide[ax0, ax1])
                    T_divide[ax0, ax1] = rxplaceholder_red[ax0, ax1] / T.Cast("float32", a * d)
    # fmt: on

    mod = LegalizeOps()(Mean)
    tvm.ir.assert_structural_equal(mod, Expected)


def test_median():
    # fmt: off
    @tvm.script.ir_module
    class Median:
        @R.function
        def main(x: R.Tensor((2, 3, 4, 5), dtype="float32")) -> R.Tuple(R.Tensor((3, 4, 5), dtype="float32"), R.Tensor((3, 4, 5), dtype="int64")):
            gv: R.Tuple(R.Tensor((3, 4, 5), dtype="float32"), R.Tensor((3, 4, 5), dtype="int64")) = R.median(x, axis=[0], keepdims=False)
            return gv

    @tvm.script.ir_module
    class Expected:
        @R.function
        def main(x: R.Tensor((2, 3, 4, 5), dtype="float32")) -> R.Tuple(R.Tensor((3, 4, 5), dtype="float32"), R.Tensor((3, 4, 5), dtype="int64")):
            gv = R.call_tir(Expected.median, (x,), out_ty=[R.Tensor((3, 4, 5), dtype="float32"), R.Tensor((3, 4, 5), dtype="int64")])
            return gv

        @T.prim_func(private=True, s_tir=True)
        def median(var_x: T.handle, T_squeeze: T.Buffer((T.int64(3), T.int64(4), T.int64(5)), "float32"), T_squeeze_1: T.Buffer((T.int64(3), T.int64(4), T.int64(5)), "int64")):
            T.func_attr({"tirx.noalias": True})
            data_buf = T.match_buffer(var_x, (T.int64(2), T.int64(3), T.int64(4), T.int64(5)), align=8)
            # with T.sblock("root"):
            T_full = T.sblock_alloc_buffer((T.int64(1), T.int64(3), T.int64(4), T.int64(5)), "int64")
            out_buf = T.sblock_alloc_buffer((T.int64(2), T.int64(3), T.int64(4), T.int64(5)), "int64", align=8)
            T_gather = T.sblock_alloc_buffer((T.int64(2), T.int64(3), T.int64(4), T.int64(5)))
            T_gather_1 = T.sblock_alloc_buffer((T.int64(1), T.int64(3), T.int64(4), T.int64(5)))
            T_gather_2 = T.sblock_alloc_buffer((T.int64(1), T.int64(3), T.int64(4), T.int64(5)), "int64")
            for ax0, ax1, ax2, ax3 in T.grid(T.int64(1), T.int64(3), T.int64(4), T.int64(5)):
                with T.sblock("T_full"):
                    v_ax0, v_ax1, v_ax2, v_ax3 = T.axis.remap("SSSS", [ax0, ax1, ax2, ax3])
                    T.reads()
                    T.writes(T_full[v_ax0, v_ax1, v_ax2, v_ax3])
                    T_full[v_ax0, v_ax1, v_ax2, v_ax3] = 0
            with T.sblock("argsort_cpu"):
                T.reads()
                T.writes()
                T.call_packed("tvm.contrib.sort.argsort", T.tvm_stack_make_array(data_buf.data,
                                                                                 T.tvm_stack_make_shape(T.int64(2), T.int64(3), T.int64(4), T.int64(5)),
                                                                                 0, 4, T.float32(0.0), T.int64(0)),
                                                          T.tvm_stack_make_array(out_buf.data,
                                                                                 T.tvm_stack_make_shape(T.int64(2), T.int64(3), T.int64(4), T.int64(5)),
                                                                                 0, 4, T.int64(0), T.int64(0)),
                                                          0, T.bool(True))
            for ax0, ax1, ax2, ax3 in T.grid(T.int64(2), T.int64(3), T.int64(4), T.int64(5)):
                with T.sblock("T_gather"):
                    v_ax0, v_ax1, v_ax2, v_ax3 = T.axis.remap("SSSS", [ax0, ax1, ax2, ax3])
                    T.reads(data_buf[out_buf[v_ax0, v_ax1, v_ax2, v_ax3], v_ax1, v_ax2, v_ax3], out_buf[v_ax0, v_ax1, v_ax2, v_ax3])
                    T.writes(T_gather[v_ax0, v_ax1, v_ax2, v_ax3])
                    T_gather[v_ax0, v_ax1, v_ax2, v_ax3] = data_buf[out_buf[v_ax0, v_ax1, v_ax2, v_ax3], v_ax1, v_ax2, v_ax3]
            for ax0, ax1, ax2, ax3 in T.grid(T.int64(1), T.int64(3), T.int64(4), T.int64(5)):
                with T.sblock("T_gather_1"):
                    v_ax0, v_ax1, v_ax2, v_ax3 = T.axis.remap("SSSS", [ax0, ax1, ax2, ax3])
                    T.reads(T_gather[T_full[v_ax0, v_ax1, v_ax2, v_ax3], v_ax1, v_ax2, v_ax3], T_full[v_ax0, v_ax1, v_ax2, v_ax3])
                    T.writes(T_gather_1[v_ax0, v_ax1, v_ax2, v_ax3])
                    T_gather_1[v_ax0, v_ax1, v_ax2, v_ax3] = T_gather[T_full[v_ax0, v_ax1, v_ax2, v_ax3], v_ax1, v_ax2, v_ax3]
            for ax0, ax1, ax2 in T.grid(T.int64(3), T.int64(4), T.int64(5)):
                with T.sblock("T_squeeze"):
                    v_ax0, v_ax1, v_ax2 = T.axis.remap("SSS", [ax0, ax1, ax2])
                    T.reads(T_gather_1[T.int64(0), v_ax0, v_ax1, v_ax2])
                    T.writes(T_squeeze[v_ax0, v_ax1, v_ax2])
                    T_squeeze[v_ax0, v_ax1, v_ax2] = T_gather_1[T.int64(0), v_ax0, v_ax1, v_ax2]
            for ax0, ax1, ax2, ax3 in T.grid(T.int64(1), T.int64(3), T.int64(4), T.int64(5)):
                with T.sblock("T_gather_2"):
                    v_ax0, v_ax1, v_ax2, v_ax3 = T.axis.remap("SSSS", [ax0, ax1, ax2, ax3])
                    T.reads(out_buf[T_full[v_ax0, v_ax1, v_ax2, v_ax3], v_ax1, v_ax2, v_ax3], T_full[v_ax0, v_ax1, v_ax2, v_ax3])
                    T.writes(T_gather_2[v_ax0, v_ax1, v_ax2, v_ax3])
                    T_gather_2[v_ax0, v_ax1, v_ax2, v_ax3] = out_buf[T_full[v_ax0, v_ax1, v_ax2, v_ax3], v_ax1, v_ax2, v_ax3]
            for ax0, ax1, ax2 in T.grid(T.int64(3), T.int64(4), T.int64(5)):
                with T.sblock("T_squeeze_1"):
                    v_ax0, v_ax1, v_ax2 = T.axis.remap("SSS", [ax0, ax1, ax2])
                    T.reads(T_gather_2[T.int64(0), v_ax0, v_ax1, v_ax2])
                    T.writes(T_squeeze_1[v_ax0, v_ax1, v_ax2])
                    T_squeeze_1[v_ax0, v_ax1, v_ax2] = T_gather_2[T.int64(0), v_ax0, v_ax1, v_ax2]
    # fmt: on

    mod = LegalizeOps()(Median)
    tvm.ir.assert_structural_equal(mod, Expected)


def test_std():
    # fmt: off
    @tvm.script.ir_module
    class Std:
        @R.function
        def main(x: R.Tensor((2, 3, 4, 5), "float32")) -> R.Tensor((), "float32"):
            gv: R.Tensor((), "float32") = R.std(x)
            return gv

    @I.ir_module(s_tir=True)
    class Expected:
        @T.prim_func(private=True, s_tir=True)
        def std(rxplaceholder: T.Buffer((T.int64(2), T.int64(3), T.int64(4), T.int64(5)), "float32"), compute: T.Buffer((), "float32")):
            T.func_attr({"tirx.noalias": True})
            # with T.sblock("root"):
            rxplaceholder_red = T.sblock_alloc_buffer((T.int64(1), T.int64(1), T.int64(1), T.int64(1)))
            T_divide = T.sblock_alloc_buffer((T.int64(1), T.int64(1), T.int64(1), T.int64(1)))
            T_subtract = T.sblock_alloc_buffer((T.int64(2), T.int64(3), T.int64(4), T.int64(5)))
            T_multiply = T.sblock_alloc_buffer((T.int64(2), T.int64(3), T.int64(4), T.int64(5)))
            T_multiply_red = T.sblock_alloc_buffer(())
            T_divide_1 = T.sblock_alloc_buffer(())
            for ax0, ax1, ax2, ax3, k0, k1, k2, k3 in T.grid(T.int64(1), T.int64(1), T.int64(1), T.int64(1), T.int64(2), T.int64(3), T.int64(4), T.int64(5)):
                with T.sblock("rxplaceholder_red"):
                    v_ax0, v_ax1, v_ax2, v_ax3, v_k0, v_k1, v_k2, v_k3 = T.axis.remap("SSSSRRRR", [ax0, ax1, ax2, ax3, k0, k1, k2, k3])
                    T.reads(rxplaceholder[v_k0, v_k1, v_k2, v_k3])
                    T.writes(rxplaceholder_red[v_ax0, v_ax1, v_ax2, v_ax3])
                    with T.init():
                        rxplaceholder_red[v_ax0, v_ax1, v_ax2, v_ax3] = T.float32(0)
                    rxplaceholder_red[v_ax0, v_ax1, v_ax2, v_ax3] = rxplaceholder_red[v_ax0, v_ax1, v_ax2, v_ax3] + rxplaceholder[v_k0, v_k1, v_k2, v_k3]
            for ax0, ax1, ax2, ax3 in T.grid(T.int64(1), T.int64(1), T.int64(1), T.int64(1)):
                with T.sblock("T_divide"):
                    v_ax0, v_ax1, v_ax2, v_ax3 = T.axis.remap("SSSS", [ax0, ax1, ax2, ax3])
                    T.reads(rxplaceholder_red[v_ax0, v_ax1, v_ax2, v_ax3])
                    T.writes(T_divide[v_ax0, v_ax1, v_ax2, v_ax3])
                    T_divide[v_ax0, v_ax1, v_ax2, v_ax3] = rxplaceholder_red[v_ax0, v_ax1, v_ax2, v_ax3] / T.float32(120.0)
            for ax0, ax1, ax2, ax3 in T.grid(T.int64(2), T.int64(3), T.int64(4), T.int64(5)):
                with T.sblock("T_subtract"):
                    v_ax0, v_ax1, v_ax2, v_ax3 = T.axis.remap("SSSS", [ax0, ax1, ax2, ax3])
                    T.reads(rxplaceholder[v_ax0, v_ax1, v_ax2, v_ax3], T_divide[T.int64(0), T.int64(0), T.int64(0), T.int64(0)])
                    T.writes(T_subtract[v_ax0, v_ax1, v_ax2, v_ax3])
                    T_subtract[v_ax0, v_ax1, v_ax2, v_ax3] = rxplaceholder[v_ax0, v_ax1, v_ax2, v_ax3] - T_divide[T.int64(0), T.int64(0), T.int64(0), T.int64(0)]
            for ax0, ax1, ax2, ax3 in T.grid(T.int64(2), T.int64(3), T.int64(4), T.int64(5)):
                with T.sblock("T_multiply"):
                    v_ax0, v_ax1, v_ax2, v_ax3 = T.axis.remap("SSSS", [ax0, ax1, ax2, ax3])
                    T.reads(T_subtract[v_ax0, v_ax1, v_ax2, v_ax3])
                    T.writes(T_multiply[v_ax0, v_ax1, v_ax2, v_ax3])
                    T_multiply[v_ax0, v_ax1, v_ax2, v_ax3] = T_subtract[v_ax0, v_ax1, v_ax2, v_ax3] * T_subtract[v_ax0, v_ax1, v_ax2, v_ax3]
            for k0, k1, k2, k3 in T.grid(T.int64(2), T.int64(3), T.int64(4), T.int64(5)):
                with T.sblock("T_multiply_red"):
                    v_k0, v_k1, v_k2, v_k3 = T.axis.remap("RRRR", [k0, k1, k2, k3])
                    T.reads(T_multiply[v_k0, v_k1, v_k2, v_k3])
                    T.writes(T_multiply_red[()])
                    with T.init():
                        T_multiply_red[()] = T.float32(0)
                    T_multiply_red[()] = T_multiply_red[()] + T_multiply[v_k0, v_k1, v_k2, v_k3]
            with T.sblock("T_divide_1"):
                vi = T.axis.spatial(1, T.int64(0))
                T.reads(T_multiply_red[()])
                T.writes(T_divide_1[()])
                T_divide_1[()] = T_multiply_red[()] / T.float32(120.0)
            with T.sblock("compute"):
                vi = T.axis.spatial(1, T.int64(0))
                T.reads(T_divide_1[()])
                T.writes(compute[()])
                compute[()] = T.sqrt(T_divide_1[()])

        @R.function
        def main(x: R.Tensor((2, 3, 4, 5), dtype="float32")) -> R.Tensor((), dtype="float32"):
            cls = Expected
            gv = R.call_tir(cls.std, (x,), out_ty=R.Tensor((), dtype="float32"))
            return gv
    # fmt: on

    mod = LegalizeOps()(Std)
    tvm.ir.assert_structural_equal(mod, Expected)


def test_std_symbolic():
    # fmt: off
    @tvm.script.ir_module
    class Std:
        @R.function
        def main(x: R.Tensor(("a", "b", "c", "d"), "float32")) -> R.Tensor((), "float32"):
            gv: R.Tensor((), "float32") = R.std(x)
            return gv

    @I.ir_module(s_tir=True)
    class Expected:
        @T.prim_func(private=True, s_tir=True)
        def std(var_rxplaceholder: T.handle, compute: T.Buffer((), "float32")):
            T.func_attr({"tirx.noalias": True})
            a, b, c, d = T.int64(), T.int64(), T.int64(), T.int64()
            rxplaceholder = T.match_buffer(var_rxplaceholder, (a, b, c, d))
            # with T.sblock("root"):
            rxplaceholder_red = T.sblock_alloc_buffer((T.int64(1), T.int64(1), T.int64(1), T.int64(1)))
            T_divide = T.sblock_alloc_buffer((T.int64(1), T.int64(1), T.int64(1), T.int64(1)))
            T_subtract = T.sblock_alloc_buffer((a, b, c, d))
            T_multiply = T.sblock_alloc_buffer((a, b, c, d))
            T_multiply_red = T.sblock_alloc_buffer(())
            T_divide_1 = T.sblock_alloc_buffer(())
            for ax0, ax1, ax2, ax3, k0, k1, k2, k3 in T.grid(T.int64(1), T.int64(1), T.int64(1), T.int64(1), a, b, c, d):
                with T.sblock("rxplaceholder_red"):
                    v_ax0, v_ax1, v_ax2, v_ax3, v_k0, v_k1, v_k2, v_k3 = T.axis.remap("SSSSRRRR", [ax0, ax1, ax2, ax3, k0, k1, k2, k3])
                    T.reads(rxplaceholder[v_k0, v_k1, v_k2, v_k3])
                    T.writes(rxplaceholder_red[v_ax0, v_ax1, v_ax2, v_ax3])
                    with T.init():
                        rxplaceholder_red[v_ax0, v_ax1, v_ax2, v_ax3] = T.float32(0)
                    rxplaceholder_red[v_ax0, v_ax1, v_ax2, v_ax3] = rxplaceholder_red[v_ax0, v_ax1, v_ax2, v_ax3] + rxplaceholder[v_k0, v_k1, v_k2, v_k3]
            for ax0, ax1, ax2, ax3 in T.grid(T.int64(1), T.int64(1), T.int64(1), T.int64(1)):
                with T.sblock("T_divide"):
                    v_ax0, v_ax1, v_ax2, v_ax3 = T.axis.remap("SSSS", [ax0, ax1, ax2, ax3])
                    T.reads(rxplaceholder_red[v_ax0, v_ax1, v_ax2, v_ax3])
                    T.writes(T_divide[v_ax0, v_ax1, v_ax2, v_ax3])
                    T_divide[v_ax0, v_ax1, v_ax2, v_ax3] = rxplaceholder_red[v_ax0, v_ax1, v_ax2, v_ax3] / T.Cast("float32", a * b * c * d)
            for ax0, ax1, ax2, ax3 in T.grid(a, b, c, d):
                with T.sblock("T_subtract"):
                    v_ax0, v_ax1, v_ax2, v_ax3 = T.axis.remap("SSSS", [ax0, ax1, ax2, ax3])
                    T.reads(rxplaceholder[v_ax0, v_ax1, v_ax2, v_ax3], T_divide[T.int64(0), T.int64(0), T.int64(0), T.int64(0)])
                    T.writes(T_subtract[v_ax0, v_ax1, v_ax2, v_ax3])
                    T_subtract[v_ax0, v_ax1, v_ax2, v_ax3] = rxplaceholder[v_ax0, v_ax1, v_ax2, v_ax3] - T_divide[T.int64(0), T.int64(0), T.int64(0), T.int64(0)]
            for ax0, ax1, ax2, ax3 in T.grid(a, b, c, d):
                with T.sblock("T_multiply"):
                    v_ax0, v_ax1, v_ax2, v_ax3 = T.axis.remap("SSSS", [ax0, ax1, ax2, ax3])
                    T.reads(T_subtract[v_ax0, v_ax1, v_ax2, v_ax3])
                    T.writes(T_multiply[v_ax0, v_ax1, v_ax2, v_ax3])
                    T_multiply[v_ax0, v_ax1, v_ax2, v_ax3] = T_subtract[v_ax0, v_ax1, v_ax2, v_ax3] * T_subtract[v_ax0, v_ax1, v_ax2, v_ax3]
            for k0, k1, k2, k3 in T.grid(a, b, c, d):
                with T.sblock("T_multiply_red"):
                    v_k0, v_k1, v_k2, v_k3 = T.axis.remap("RRRR", [k0, k1, k2, k3])
                    T.reads(T_multiply[v_k0, v_k1, v_k2, v_k3])
                    T.writes(T_multiply_red[()])
                    with T.init():
                        T_multiply_red[()] = T.float32(0)
                    T_multiply_red[()] = T_multiply_red[()] + T_multiply[v_k0, v_k1, v_k2, v_k3]
            with T.sblock("T_divide_1"):
                vi = T.axis.spatial(1, T.int64(0))
                T.reads(T_multiply_red[()])
                T.writes(T_divide_1[()])
                T_divide_1[()] = T_multiply_red[()] / T.Cast("float32", a * b * c * d)
            with T.sblock("compute"):
                vi = T.axis.spatial(1, T.int64(0))
                T.reads(T_divide_1[()])
                T.writes(compute[()])
                compute[()] = T.sqrt(T_divide_1[()])

        @R.function
        def main(x: R.Tensor(("a", "b", "c", "d"), dtype="float32")) -> R.Tensor((), dtype="float32"):
            a = T.int64()
            b = T.int64()
            c = T.int64()
            d = T.int64()
            cls = Expected
            gv = R.call_tir(cls.std, (x,), out_ty=R.Tensor((), dtype="float32"))
            return gv
    # fmt: on

    mod = LegalizeOps()(Std)
    tvm.ir.assert_structural_equal(mod, Expected)


def test_variance():
    # fmt: off
    @tvm.script.ir_module
    class Variance:
        @R.function
        def main(x: R.Tensor((2, 3, 4, 5), "float32")) -> R.Tensor((1, 3, 4, 1), "float32"):
            gv: R.Tensor((1, 3, 4, 1), "float32") = R.variance(x, [0, 3], keepdims=True)
            return gv

    @tvm.script.ir_module
    class Expected:
        @R.function
        def main(x: R.Tensor((2, 3, 4, 5), dtype="float32")) -> R.Tensor((1, 3, 4, 1), dtype="float32"):
            gv = R.call_tir(Expected.variance, (x,), R.Tensor((1, 3, 4, 1), dtype="float32"))
            return gv

        @T.prim_func(private=True, s_tir=True)
        def variance(rxplaceholder: T.Buffer((T.int64(2), T.int64(3), T.int64(4), T.int64(5)), "float32"), T_divide: T.Buffer((T.int64(1), T.int64(3), T.int64(4), T.int64(1)), "float32")):
            T.func_attr({"tirx.noalias": True})
            rxplaceholder_red = T.sblock_alloc_buffer([T.int64(1), T.int64(3), T.int64(4), T.int64(1)], dtype="float32")
            T_divide_1 = T.sblock_alloc_buffer([T.int64(1), T.int64(3), T.int64(4), T.int64(1)], dtype="float32")
            T_subtract = T.sblock_alloc_buffer([T.int64(2), T.int64(3), T.int64(4), T.int64(5)], dtype="float32")
            T_multiply = T.sblock_alloc_buffer([T.int64(2), T.int64(3), T.int64(4), T.int64(5)], dtype="float32")
            T_multiply_red = T.sblock_alloc_buffer([T.int64(1), T.int64(3), T.int64(4), T.int64(1)], dtype="float32")
            for i0, i1, i2, i3, i4, i5 in T.grid(T.int64(1), T.int64(3), T.int64(4), T.int64(1), T.int64(2), T.int64(5)):
                with T.sblock("rxplaceholder_red"):
                    ax0, ax1, ax2, ax3, k0, k3 = T.axis.remap("SSSSRR", [i0, i1, i2, i3, i4, i5])
                    T.reads(rxplaceholder[k0, ax1, ax2, k3])
                    T.writes(rxplaceholder_red[ax0, ax1, ax2, ax3])
                    with T.init():
                        rxplaceholder_red[ax0, ax1, ax2, ax3] = T.float32(0)
                    rxplaceholder_red[ax0, ax1, ax2, ax3] = rxplaceholder_red[ax0, ax1, ax2, ax3] + rxplaceholder[k0, ax1, ax2, k3]
            for i0, i1, i2, i3 in T.grid(T.int64(1), T.int64(3), T.int64(4), T.int64(1)):
                with T.sblock("T_divide"):
                    ax0, ax1, ax2, ax3 = T.axis.remap("SSSS", [i0, i1, i2, i3])
                    T.reads(rxplaceholder_red[ax0, ax1, ax2, ax3])
                    T.writes(T_divide_1[ax0, ax1, ax2, ax3])
                    T_divide_1[ax0, ax1, ax2, ax3] = rxplaceholder_red[ax0, ax1, ax2, ax3] / T.float32(10.0)
            for i0, i1, i2, i3 in T.grid(T.int64(2), T.int64(3), T.int64(4), T.int64(5)):
                with T.sblock("T_subtract"):
                    ax0, ax1, ax2, ax3 = T.axis.remap("SSSS", [i0, i1, i2, i3])
                    T.reads(rxplaceholder[ax0, ax1, ax2, ax3], T_divide_1[T.int64(0), ax1, ax2, T.int64(0)])
                    T.writes(T_subtract[ax0, ax1, ax2, ax3])
                    T_subtract[ax0, ax1, ax2, ax3] = rxplaceholder[ax0, ax1, ax2, ax3] - T_divide_1[T.int64(0), ax1, ax2, T.int64(0)]
            for i0, i1, i2, i3 in T.grid(T.int64(2), T.int64(3), T.int64(4), T.int64(5)):
                with T.sblock("T_multiply"):
                    ax0, ax1, ax2, ax3 = T.axis.remap("SSSS", [i0, i1, i2, i3])
                    T.reads(T_subtract[ax0, ax1, ax2, ax3])
                    T.writes(T_multiply[ax0, ax1, ax2, ax3])
                    T_multiply[ax0, ax1, ax2, ax3] = T_subtract[ax0, ax1, ax2, ax3] * T_subtract[ax0, ax1, ax2, ax3]
            for i0, i1, i2, i3, i4, i5 in T.grid(T.int64(1), T.int64(3), T.int64(4), T.int64(1), T.int64(2), T.int64(5)):
                with T.sblock("T_multiply_red"):
                    ax0, ax1, ax2, ax3, k0, k3 = T.axis.remap("SSSSRR", [i0, i1, i2, i3, i4, i5])
                    T.reads(T_multiply[k0, ax1, ax2, k3])
                    T.writes(T_multiply_red[ax0, ax1, ax2, ax3])
                    with T.init():
                        T_multiply_red[ax0, ax1, ax2, ax3] = T.float32(0)
                    T_multiply_red[ax0, ax1, ax2, ax3] = T_multiply_red[ax0, ax1, ax2, ax3] + T_multiply[k0, ax1, ax2, k3]
            for i0, i1, i2, i3 in T.grid(T.int64(1), T.int64(3), T.int64(4), T.int64(1)):
                with T.sblock("T_divide_1"):
                    ax0, ax1, ax2, ax3 = T.axis.remap("SSSS", [i0, i1, i2, i3])
                    T.reads(T_multiply_red[ax0, ax1, ax2, ax3])
                    T.writes(T_divide[ax0, ax1, ax2, ax3])
                    T_divide[ax0, ax1, ax2, ax3] = T_multiply_red[ax0, ax1, ax2, ax3] / T.float32(10)
    # fmt: on

    mod = LegalizeOps()(Variance)
    tvm.ir.assert_structural_equal(mod, Expected)


def test_variance_symbolic():
    # fmt: off
    @tvm.script.ir_module
    class Variance:
        @R.function
        def main(x: R.Tensor(("a", "b", "c", "d"), "float32")) -> R.Tensor((1, "b", "c", 1), "float32"):
            b = T.int64()
            c = T.int64()
            gv: R.Tensor((1, b, c, 1), "float32") = R.variance(x, [0, 3], keepdims=True)
            return gv

    @tvm.script.ir_module
    class Expected:
        @R.function
        def main(x: R.Tensor(("a", "b", "c", "d"), "float32")) -> R.Tensor((1, "b", "c", 1), "float32"):
            b = T.int64()
            c = T.int64()
            gv = R.call_tir(Expected.variance, (x,), R.Tensor((1, b, c, 1), dtype="float32"))
            return gv

        @T.prim_func(private=True, s_tir=True)
        def variance(var_rxplaceholder: T.handle, var_T_divide: T.handle):
            T.func_attr({"tirx.noalias": True})
            a = T.int64()
            b = T.int64()
            c = T.int64()
            d = T.int64()
            rxplaceholder = T.match_buffer(var_rxplaceholder, [a, b, c, d], dtype="float32")
            T_divide = T.match_buffer(var_T_divide, [T.int64(1), b, c, T.int64(1)], dtype="float32")
            rxplaceholder_red = T.sblock_alloc_buffer([T.int64(1), b, c, T.int64(1)], dtype="float32")
            T_divide_1 = T.sblock_alloc_buffer([T.int64(1), b, c, T.int64(1)], dtype="float32")
            T_subtract = T.sblock_alloc_buffer([a, b, c, d], dtype="float32")
            T_multiply = T.sblock_alloc_buffer([a, b, c, d], dtype="float32")
            T_multiply_red = T.sblock_alloc_buffer([T.int64(1), b, c, T.int64(1)], dtype="float32")
            for i0, i1, i2, i3, i4, i5 in T.grid(T.int64(1), b, c, T.int64(1), a, d):
                with T.sblock("rxplaceholder_red"):
                    ax0, ax1, ax2, ax3, k0, k3 = T.axis.remap("SSSSRR", [i0, i1, i2, i3, i4, i5])
                    T.reads(rxplaceholder[k0, ax1, ax2, k3])
                    T.writes(rxplaceholder_red[ax0, ax1, ax2, ax3])
                    with T.init():
                        rxplaceholder_red[ax0, ax1, ax2, ax3] = T.float32(0)
                    rxplaceholder_red[ax0, ax1, ax2, ax3] = rxplaceholder_red[ax0, ax1, ax2, ax3] + rxplaceholder[k0, ax1, ax2, k3]
            for i0, i1, i2, i3 in T.grid(T.int64(1), b, c, T.int64(1)):
                with T.sblock("T_divide"):
                    ax0, ax1, ax2, ax3 = T.axis.remap("SSSS", [i0, i1, i2, i3])
                    T.reads(rxplaceholder_red[ax0, ax1, ax2, ax3])
                    T.writes(T_divide_1[ax0, ax1, ax2, ax3])
                    T_divide_1[ax0, ax1, ax2, ax3] = rxplaceholder_red[ax0, ax1, ax2, ax3] / T.Cast("float32", a * d)
            for i0, i1, i2, i3 in T.grid(a, b, c, d):
                with T.sblock("T_subtract"):
                    ax0, ax1, ax2, ax3 = T.axis.remap("SSSS", [i0, i1, i2, i3])
                    T.reads(rxplaceholder[ax0, ax1, ax2, ax3], T_divide_1[T.int64(0), ax1, ax2, T.int64(0)])
                    T.writes(T_subtract[ax0, ax1, ax2, ax3])
                    T_subtract[ax0, ax1, ax2, ax3] = rxplaceholder[ax0, ax1, ax2, ax3] - T_divide_1[T.int64(0), ax1, ax2, T.int64(0)]
            for i0, i1, i2, i3 in T.grid(a, b, c, d):
                with T.sblock("T_multiply"):
                    ax0, ax1, ax2, ax3 = T.axis.remap("SSSS", [i0, i1, i2, i3])
                    T.reads(T_subtract[ax0, ax1, ax2, ax3])
                    T.writes(T_multiply[ax0, ax1, ax2, ax3])
                    T_multiply[ax0, ax1, ax2, ax3] = T_subtract[ax0, ax1, ax2, ax3] * T_subtract[ax0, ax1, ax2, ax3]
            for i0, i1, i2, i3, i4, i5 in T.grid(T.int64(1), b, c, T.int64(1), a, d):
                with T.sblock("T_multiply_red"):
                    ax0, ax1, ax2, ax3, k0, k3 = T.axis.remap("SSSSRR", [i0, i1, i2, i3, i4, i5])
                    T.reads(T_multiply[k0, ax1, ax2, k3])
                    T.writes(T_multiply_red[ax0, ax1, ax2, ax3])
                    with T.init():
                        T_multiply_red[ax0, ax1, ax2, ax3] = T.float32(0)
                    T_multiply_red[ax0, ax1, ax2, ax3] = T_multiply_red[ax0, ax1, ax2, ax3] + T_multiply[k0, ax1, ax2, k3]
            for i0, i1, i2, i3 in T.grid(T.int64(1), b, c, T.int64(1)):
                with T.sblock("T_divide_1"):
                    ax0, ax1, ax2, ax3 = T.axis.remap("SSSS", [i0, i1, i2, i3])
                    T.reads(T_multiply_red[ax0, ax1, ax2, ax3])
                    T.writes(T_divide[ax0, ax1, ax2, ax3])
                    T_divide[ax0, ax1, ax2, ax3] = T_multiply_red[ax0, ax1, ax2, ax3] / T.Cast("float32", a * d)
    # fmt: on

    mod = LegalizeOps()(Variance)
    tvm.ir.assert_structural_equal(mod, Expected)


def test_variance_no_keepdims():
    # fmt: off
    @tvm.script.ir_module
    class Variance:
        @R.function
        def main(x: R.Tensor((2, 3, 4, 5), "float32")) -> R.Tensor((3, 4), "float32"):
            gv: R.Tensor((3, 4), "float32") = R.variance(x, [0, 3], keepdims=False)
            return gv

    @I.ir_module(s_tir=True)
    class Expected:
        @T.prim_func(private=True, s_tir=True)
        def variance(rxplaceholder: T.Buffer((T.int64(2), T.int64(3), T.int64(4), T.int64(5)), "float32"), T_divide: T.Buffer((T.int64(3), T.int64(4)), "float32")):
            T.func_attr({"tirx.noalias": True})
            # with T.sblock("root"):
            rxplaceholder_red = T.sblock_alloc_buffer((T.int64(1), T.int64(3), T.int64(4), T.int64(1)))
            T_divide_1 = T.sblock_alloc_buffer((T.int64(1), T.int64(3), T.int64(4), T.int64(1)))
            T_subtract = T.sblock_alloc_buffer((T.int64(2), T.int64(3), T.int64(4), T.int64(5)))
            T_multiply = T.sblock_alloc_buffer((T.int64(2), T.int64(3), T.int64(4), T.int64(5)))
            T_multiply_red = T.sblock_alloc_buffer((T.int64(3), T.int64(4)))
            for ax0, ax1, ax2, ax3, k0, k3 in T.grid(T.int64(1), T.int64(3), T.int64(4), T.int64(1), T.int64(2), T.int64(5)):
                with T.sblock("rxplaceholder_red"):
                    v_ax0, v_ax1, v_ax2, v_ax3, v_k0, v_k3 = T.axis.remap("SSSSRR", [ax0, ax1, ax2, ax3, k0, k3])
                    T.reads(rxplaceholder[v_k0, v_ax1, v_ax2, v_k3])
                    T.writes(rxplaceholder_red[v_ax0, v_ax1, v_ax2, v_ax3])
                    with T.init():
                        rxplaceholder_red[v_ax0, v_ax1, v_ax2, v_ax3] = T.float32(0)
                    rxplaceholder_red[v_ax0, v_ax1, v_ax2, v_ax3] = rxplaceholder_red[v_ax0, v_ax1, v_ax2, v_ax3] + rxplaceholder[v_k0, v_ax1, v_ax2, v_k3]
            for ax0, ax1, ax2, ax3 in T.grid(T.int64(1), T.int64(3), T.int64(4), T.int64(1)):
                with T.sblock("T_divide"):
                    v_ax0, v_ax1, v_ax2, v_ax3 = T.axis.remap("SSSS", [ax0, ax1, ax2, ax3])
                    T.reads(rxplaceholder_red[v_ax0, v_ax1, v_ax2, v_ax3])
                    T.writes(T_divide_1[v_ax0, v_ax1, v_ax2, v_ax3])
                    T_divide_1[v_ax0, v_ax1, v_ax2, v_ax3] = rxplaceholder_red[v_ax0, v_ax1, v_ax2, v_ax3] / T.float32(10)
            for ax0, ax1, ax2, ax3 in T.grid(T.int64(2), T.int64(3), T.int64(4), T.int64(5)):
                with T.sblock("T_subtract"):
                    v_ax0, v_ax1, v_ax2, v_ax3 = T.axis.remap("SSSS", [ax0, ax1, ax2, ax3])
                    T.reads(rxplaceholder[v_ax0, v_ax1, v_ax2, v_ax3], T_divide_1[T.int64(0), v_ax1, v_ax2, T.int64(0)])
                    T.writes(T_subtract[v_ax0, v_ax1, v_ax2, v_ax3])
                    T_subtract[v_ax0, v_ax1, v_ax2, v_ax3] = rxplaceholder[v_ax0, v_ax1, v_ax2, v_ax3] - T_divide_1[T.int64(0), v_ax1, v_ax2, T.int64(0)]
            for ax0, ax1, ax2, ax3 in T.grid(T.int64(2), T.int64(3), T.int64(4), T.int64(5)):
                with T.sblock("T_multiply"):
                    v_ax0, v_ax1, v_ax2, v_ax3 = T.axis.remap("SSSS", [ax0, ax1, ax2, ax3])
                    T.reads(T_subtract[v_ax0, v_ax1, v_ax2, v_ax3])
                    T.writes(T_multiply[v_ax0, v_ax1, v_ax2, v_ax3])
                    T_multiply[v_ax0, v_ax1, v_ax2, v_ax3] = T_subtract[v_ax0, v_ax1, v_ax2, v_ax3] * T_subtract[v_ax0, v_ax1, v_ax2, v_ax3]
            for ax0, ax1, k0, k3 in T.grid(T.int64(3), T.int64(4), T.int64(2), T.int64(5)):
                with T.sblock("T_multiply_red"):
                    v_ax0, v_ax1, v_k0, v_k3 = T.axis.remap("SSRR", [ax0, ax1, k0, k3])
                    T.reads(T_multiply[v_k0, v_ax0, v_ax1, v_k3])
                    T.writes(T_multiply_red[v_ax0, v_ax1])
                    with T.init():
                        T_multiply_red[v_ax0, v_ax1] = T.float32(0)
                    T_multiply_red[v_ax0, v_ax1] = T_multiply_red[v_ax0, v_ax1] + T_multiply[v_k0, v_ax0, v_ax1, v_k3]
            for ax0, ax1 in T.grid(T.int64(3), T.int64(4)):
                with T.sblock("T_divide_1"):
                    v_ax0, v_ax1 = T.axis.remap("SS", [ax0, ax1])
                    T.reads(T_multiply_red[v_ax0, v_ax1])
                    T.writes(T_divide[v_ax0, v_ax1])
                    T_divide[v_ax0, v_ax1] = T_multiply_red[v_ax0, v_ax1] / T.float32(10)

        @R.function
        def main(x: R.Tensor((2, 3, 4, 5), dtype="float32")) -> R.Tensor((3, 4), dtype="float32"):
            cls = Expected
            gv = R.call_tir(cls.variance, (x,), out_ty=R.Tensor((3, 4), dtype="float32"))
            return gv
    # fmt: on

    mod = LegalizeOps()(Variance)
    tvm.ir.assert_structural_equal(mod, Expected)


def test_max_zero_dim():
    # Reducing a 0-D (scalar) tensor is the identity; it must legalize, not crash.
    # Regression test for https://github.com/apache/tvm/issues/19676
    # fmt: off
    @tvm.script.ir_module
    class Max:
        @R.function
        def main(x: R.Tensor((), "float32")) -> R.Tensor((), "float32"):
            gv: R.Tensor((), "float32") = R.max(x)
            return gv

    @tvm.script.ir_module
    class Expected:
        @R.function
        def main(x: R.Tensor((), dtype="float32")) -> R.Tensor((), dtype="float32"):
            gv = R.call_tir(Expected.max, (x,), R.Tensor((), dtype="float32"))
            return gv

        @T.prim_func(private=True, s_tir=True)
        def max(x: T.Buffer((), "float32"), x_red: T.Buffer((), "float32")):
            T.func_attr({"tirx.noalias": True})
            with T.sblock("x_red"):
                vi = T.axis.spatial(1, T.int64(0))
                T.reads(x[()])
                T.writes(x_red[()])
                x_red[()] = x[()]
    # fmt: on

    mod = LegalizeOps()(Max)
    tvm.ir.assert_structural_equal(mod, Expected)


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