# 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.
"""Backend compiler related feature registration"""
# pylint: disable=invalid-name,unused-argument, len-as-condition, too-many-nested-blocks,
# pylint: disable=too-many-local-variables, too-many-arguments, no-else-return

from __future__ import absolute_import

import tvm
from tvm import te, topi
from tvm.runtime import convert
from tvm.te.hybrid import script
from tvm.topi.utils import get_const_int, get_const_tuple

from . import op as _reg
from . import strategy
from ._tensor import elemwise_shape_func
from .op import OpPattern

_reg.register_broadcast_schedule("broadcast_to")
_reg.register_broadcast_schedule("broadcast_to_like")
_reg.register_broadcast_schedule("expand_dims")
_reg.register_broadcast_schedule("repeat")
_reg.register_broadcast_schedule("tile")
_reg.register_broadcast_schedule("where")
_reg.register_injective_schedule("squeeze")
_reg.register_injective_schedule("reshape")
_reg.register_injective_schedule("reshape_like")
_reg.register_injective_schedule("full")
_reg.register_injective_schedule("full_like")
_reg.register_injective_schedule("arange")
_reg.register_injective_schedule("meshgrid")
_reg.register_injective_schedule("reverse")
_reg.register_injective_schedule("reverse_sequence")
_reg.register_injective_schedule("cast")
_reg.register_injective_schedule("cast_like")
_reg.register_injective_schedule("reinterpret")
_reg.register_injective_schedule("strided_slice")
_reg.register_injective_schedule("slice_like")
_reg.register_injective_schedule("split")
_reg.register_injective_schedule("take")
_reg.register_injective_schedule("stack")
_reg.register_injective_schedule("contrib_reverse_reshape")
_reg.register_injective_schedule("gather")
_reg.register_injective_schedule("gather_nd")
_reg.register_injective_schedule("sequence_mask")
_reg.register_injective_schedule("one_hot")
_reg.register_reduce_schedule("collapse_sum_like")
_reg.register_reduce_schedule("collapse_sum_to")
_reg.register_injective_schedule("unravel_index")
_reg.register_injective_schedule("sparse_to_dense")
_reg.register_injective_schedule("matrix_set_diag")
_reg.register_injective_schedule("adv_index")


# concatenate
@_reg.register_compute("concatenate")
def compute_concat(attrs, inputs, output_type):
    return [topi.concatenate(inputs, attrs.axis)]


_reg.register_strategy("concatenate", strategy.concatenate_strategy)

# sliding_window
@_reg.register_compute("sliding_window")
def compute_sliding_window(attrs, inputs, output_type):
    """Compute definition of sliding_window"""
    return [topi.sliding_window(inputs[0], attrs.axis, attrs.window_shape, attrs.strides)]


_reg.register_strategy("sliding_window", strategy.sliding_window_strategy)

# strided_set
@_reg.register_compute("strided_set")
def compute_strided_set(attrs, inputs, output_type):
    """Compute definition of strided_set"""
    return [topi.strided_set(inputs[0], inputs[1], inputs[2], inputs[3], inputs[4])]


_reg.register_injective_schedule("strided_set")

# layout_transform
_reg.register_injective_schedule("layout_transform")
_reg.register_pattern("layout_transform", OpPattern.INJECTIVE)
_reg.register_injective_schedule("auto_scheduler_layout_transform")
_reg.register_pattern("auto_scheduler_layout_transform", OpPattern.INJECTIVE)
_reg.register_injective_schedule("meta_schedule_layout_transform")
_reg.register_pattern("meta_schedule_layout_transform", OpPattern.INJECTIVE)

# argwhere
_reg.register_strategy("argwhere", strategy.argwhere_strategy)

# scatter
@_reg.register_compute("scatter")
def compute_scatter(attrs, inputs, output_type):
    """Compute definition of scatter"""
    return [topi.scatter(inputs[0], inputs[1], inputs[2], attrs.axis)]


_reg.register_strategy("scatter", strategy.scatter_strategy)

# sparse_fill_empty_rows
@_reg.register_compute("sparse_fill_empty_rows")
def compute_sparse_fill_empty_rows(attrs, inputs, output_type):
    """Compute definition of sparse_fill_empty_rows"""

    return topi.sparse_fill_empty_rows(
        inputs[0],
        inputs[1],
        inputs[2],
        inputs[3],
        output_type.fields[0].shape,
        output_type.fields[1].shape,
        output_type.fields[2].shape,
    )


_reg.register_strategy("sparse_fill_empty_rows", strategy.sparse_fill_empty_rows_strategy)

# sparse_reshape
@_reg.register_compute("sparse_reshape")
def compute_reshape(attrs, inputs, output_type):
    """Compute definition of sparse_reshape"""

    return topi.sparse_reshape(
        inputs[0],
        inputs[1],
        inputs[2],
        output_type.fields[0].shape,
        output_type.fields[1].shape,
    )


_reg.register_strategy("sparse_reshape", strategy.sparse_reshape_strategy)

# stft
@_reg.register_compute("stft")
def compute_stft(attrs, inputs, output_type):
    """Compute definition of stft"""
    return topi.stft(
        inputs[0],
        attrs.n_fft,
        attrs.hop_length,
        attrs.win_length,
        attrs.window,
        attrs.normalized,
        attrs.onesided,
        output_type.shape,
    )


_reg.register_strategy("stft", strategy.stft_strategy)


@script
def _stft_shape_func(data, n_fft, hop_length, onesided):
    output_shape = output_tensor((4,), "int64")
    output_shape[0] = int64(data.shape[0])
    if onesided:
        output_shape[1] = int64(int64(n_fft) // int64(2)) + int64(1)
    else:
        output_shape[1] = int64(n_fft)
    output_shape[2] = int64(int64(data.shape[1] - n_fft) // int64(hop_length)) + int64(1)
    output_shape[3] = int64(2)
    return output_shape


@_reg.register_shape_func("stft", True)
def stft_shape_func(attrs, inputs, _):
    """
    Shape func for stft.
    """
    return [
        _stft_shape_func(
            inputs[0], convert(attrs.n_fft), convert(attrs.hop_length), convert(attrs.onesided)
        )
    ]


# trilu
_reg.register_strategy("trilu", strategy.trilu_strategy)


# scatter_add
@_reg.register_compute("scatter_add")
def compute_scatter_add(attrs, inputs, output_type):
    """Compute definition of scatter_add"""
    return [topi.scatter_add(inputs[0], inputs[1], inputs[2], attrs.axis)]


_reg.register_strategy("scatter_add", strategy.scatter_add_strategy)

# scatter_nd
@_reg.register_compute("scatter_nd")
def compute_scatter_nd(attrs, inputs, output_type):
    """Compute definition of scatter_nd"""
    return [topi.scatter_nd(inputs[0], inputs[1], inputs[2], attrs.mode)]


_reg.register_strategy("scatter_nd", strategy.scatter_nd_strategy)

# cumsum
@_reg.register_compute("cumsum")
def compute_cumsum(attrs, inputs, output_type):
    """Compute definition of cumsum"""
    return [topi.cumsum(inputs[0], attrs.axis, attrs.dtype, attrs.exclusive)]


_reg.register_strategy("cumsum", strategy.cumsum_strategy)
_reg.register_shape_func("cumsum", False, elemwise_shape_func)

# cumprod
@_reg.register_compute("cumprod")
def compute_cumprod(attrs, inputs, output_type):
    """Compute definition of cumprod"""
    return [topi.cumprod(inputs[0], attrs.axis, attrs.dtype, attrs.exclusive)]


_reg.register_strategy("cumprod", strategy.cumprod_strategy)
_reg.register_shape_func("cumprod", False, elemwise_shape_func)


@_reg.register_compute("unique")
def compute_unique(attrs, inputs, output_type):
    """Compute definition of unique"""
    return topi.unique(inputs[0], attrs.sorted, attrs.return_counts)


_reg.register_strategy("unique", strategy.unique_strategy)

# invert_permutation
_reg.register_strategy("invert_permutation", strategy.invert_permutation_strategy)
_reg.register_shape_func("invert_permutation", False, elemwise_shape_func)


#####################
#  Shape functions  #
#####################


@script
def _arange_shape_func(start, stop, step):
    out = output_tensor((1,), "int64")
    if step[()] < 0:
        out[0] = int64(ceil_div((int64(start[()]) - int64(stop[()])), int64(-step[()])))
    else:
        out[0] = int64(ceil_div((int64(stop[()]) - int64(start[()])), int64(step[()])))
    return out


@_reg.register_shape_func("arange", True)
def arange_shape_func(attrs, inputs, _):
    """
    Shape func for arange
    """
    return [_arange_shape_func(*inputs)]


@script
def _strided_slice_shape_func_input_shape(data_shape, begin, end, strides, slice_mode):
    ndim = len(data_shape)
    out = output_tensor((ndim,), "int64")
    for i in const_range(ndim):
        dim_size = int64(data_shape[i])
        cbegin = int64(0)
        cend = dim_size
        cstride = int64(1)

        if len(strides) > i:
            cstride = int64(strides[i])

        if len(begin) > i:
            cbegin = int64(begin[i])
        elif cstride < 0:
            cbegin = dim_size

        if len(end) <= i:
            if cstride < 0:
                cend = int64(0)
        elif slice_mode != 0:
            cstride = int64(1)
            if end[i] < 0:
                cend = dim_size
            else:
                cend = cbegin + int64(end[i])
        else:
            if end[i] > data_shape[i]:
                cend = dim_size
            else:
                cend = int64(end[i])

        assert cstride != 0, "Strides can't be zero."

        if cbegin < 0:
            cbegin += dim_size
        if cend < 0:
            cend += dim_size

        if cstride < 0:
            if cend < 0:
                cend = int64(-1)
            if cbegin > dim_size - 1:
                cbegin = dim_size - 1
            slice_range = cbegin - cend
            step = -cstride
        else:
            slice_range = cend - cbegin
            step = cstride
        out[i] = int64(ceil_div(slice_range, step))
    return out


@script
def _strided_slice_shape_func_with_axes(data_shape, begin, end, strides, slice_mode, axes):
    ndim = data_shape.shape[0]
    out = output_tensor((ndim,), "int64")
    for i in const_range(ndim):
        out[i] = data_shape[i]

    for i in const_range(len(axes)):
        dim_size = int64(data_shape[axes[i]])
        cbegin = int64(0)
        cend = dim_size
        cstride = int64(1)

        if len(strides) > i:
            cstride = int64(strides[i])

        if len(begin) > i:
            cbegin = int64(begin[i])
        elif cstride < 0:
            cbegin = dim_size

        if len(end) <= i:
            cend = dim_size
        elif slice_mode != 0:
            cstride = int64(1)
            if end[i] < 0:
                cend = dim_size
            else:
                cend = cbegin + int64(end[i])
        else:
            if end[i] > data_shape[axes[i]]:
                cend = dim_size
            else:
                cend = int64(end[i])

        assert cstride != 0, "Strides can't be zero."

        if cbegin < 0:
            cbegin += dim_size
        if cend < 0:
            cend += dim_size

        if cstride < 0:
            if cend < 0:
                cend = int64(-1)
            if cbegin > dim_size - 1:
                cbegin = dim_size - 1
            slice_range = cbegin - cend
            step = -cstride
        else:
            slice_range = cend - cbegin
            step = cstride

        out[axes[i]] = int64(ceil_div(slice_range, step))
    return out


@_reg.register_shape_func("strided_slice", False)
def strided_slice_shape_func(attrs, inputs, _):
    """
    Shape func for strided_slice
    """
    slice_mode = convert(0 if attrs.slice_mode == "end" else 1)
    if attrs.axes is None:
        return [
            _strided_slice_shape_func_input_shape(
                inputs[0], attrs.begin, attrs.end, attrs.strides, slice_mode
            )
        ]
    return [
        _strided_slice_shape_func_with_axes(
            inputs[0], attrs.begin, attrs.end, attrs.strides, slice_mode, attrs.axes
        )
    ]


@script
def _one_hot_shape_func(indices_shape, depth, axis):
    in_ndim = indices_shape.shape[0]
    out_ndim = in_ndim + 1
    true_axis = in_ndim if axis == -1 else axis
    indices_i = 0
    out = output_tensor((out_ndim,), "int64")
    for i in range(out_ndim):
        if i == true_axis:
            out[i] = int64(depth)
        else:
            out[i] = int64(indices_shape[indices_i])
            indices_i += 1
    return out


@_reg.register_shape_func("one_hot", False)
def one_hot_shape_func(attrs, inputs, _):
    """
    Shape func for one_hot
    """
    shape_func = [_one_hot_shape_func(inputs[0], convert(attrs.depth), convert(attrs.axis))]
    return shape_func


@script
def _concatenate_shape_func(inputs, axis):
    ndim = inputs[0].shape[0]
    out = output_tensor((ndim,), "int64")
    for i in const_range(ndim):
        if i != axis:
            out[i] = inputs[0][i]
            for j in const_range(1, len(inputs)):
                assert out[i] == inputs[j][i], "Dims mismatch in the inputs of concatenate."
        else:
            out[i] = int64(0)
            for j in const_range(len(inputs)):
                out[i] += inputs[j][i]
    return out


@_reg.register_shape_func("concatenate", False)
def concatenate_shape_func(attrs, inputs, _):
    axis = get_const_int(attrs.axis)
    if axis < 0:
        axis += inputs[0].shape[0]
    return [_concatenate_shape_func(inputs, convert(axis))]


@script
def _reshape_shape_func_input_shape(data_shape, newshape, ndim, allowzero):
    out = output_tensor((ndim,), "int64")
    src_idx = 0
    dst_idx = 0
    infer_idx = -1
    copy = False
    skip = 0
    for i in const_range(len(newshape)):
        if skip > 0:
            skip -= 1
        elif newshape[i] > 0:
            out[dst_idx] = int64(newshape[i])
            src_idx += 1
            dst_idx += 1
        elif newshape[i] == 0:
            if allowzero:
                out[dst_idx] = int64(newshape[i])
            else:
                out[dst_idx] = data_shape[src_idx]
            src_idx += 1
            dst_idx += 1
        elif newshape[i] == -1:
            assert infer_idx < 0, "One and only one dim can be inferred"
            out[dst_idx] = int64(1)
            infer_idx = i
            src_idx += 1
            dst_idx += 1
        elif newshape[i] == -2:
            copy = True
        elif newshape[i] == -3:
            assert data_shape.shape[0] - src_idx > 1, "Not enough dims in input shape for -3"
            out[dst_idx] = data_shape[src_idx] * data_shape[src_idx + 1]
            src_idx += 2
            dst_idx += 1
        elif newshape[i] == -4:
            assert len(newshape) - i > 2, "Not enough dims in new shape for -4"
            if newshape[i + 1] == -1:
                assert newshape[i + 2] != -1, "Split dims cannot both be -1."
                out[dst_idx] = data_shape[src_idx] // int64(newshape[i + 2])
                out[dst_idx + 1] = int64(newshape[i + 2])
            else:
                out[dst_idx] = int64(newshape[i + 1])
                if newshape[i + 2] == -1:
                    out[dst_idx + 1] = data_shape[src_idx] // int64(newshape[i + 1])
                else:
                    out[dst_idx + 1] = int64(newshape[i + 2])
            assert (
                data_shape[src_idx] == out[dst_idx] * out[dst_idx + 1]
            ), "Product of split dims doesn't match to input dim"
            src_idx += 1
            dst_idx += 2
            skip = 2
        else:
            assert False, "Invalid special values in new shape"
    if len(data_shape.shape) > 0:
        # if data is not constant, we can then handle -1 and -2
        if copy:
            for i in range(src_idx, data_shape.shape[0]):
                out[dst_idx] = data_shape[i]
                dst_idx += 1
        if infer_idx >= 0:
            old_size = int64(1)
            for i in const_range(data_shape.shape[0]):
                old_size *= data_shape[i]
            new_size = int64(1)
            for i in const_range(out.shape[0]):
                new_size *= out[i]
            out[infer_idx] = old_size // new_size
    return out


@_reg.register_shape_func("reshape", False)
def reshape_shape_func(attrs, inputs, out_ndims):
    newshape = get_const_tuple(attrs.newshape)
    allowzero = attrs.allowzero
    return [
        _reshape_shape_func_input_shape(
            inputs[0], convert(newshape), out_ndims[0], convert(allowzero)
        )
    ]


@script
def _take_no_axis_shape_func(indices_shape, out_ndim):
    out = output_tensor((out_ndim,), "int64")
    for i in const_range(out_ndim):
        out[i] = indices_shape[i]
    return out


@script
def _take_with_axis_shape_func(data_shape, indices_shape, axis, batch_dims, out_ndim):
    out = output_tensor((out_ndim,), "int64")
    for i in const_range(axis):
        out[i] = data_shape[i]
    if len(indices_shape.shape) == 0:
        # indices is constant
        for i in const_range(axis + 1, len(data_shape)):
            out[i - 1] = data_shape[i]
    else:
        for i in const_range(len(indices_shape) - batch_dims):
            out[axis + i] = indices_shape[i + batch_dims]
        for i in const_range(axis + 1, len(data_shape)):
            out[len(indices_shape) + i - 1 - batch_dims] = data_shape[i]
    return out


@_reg.register_shape_func("take", False)
def take_shape_func(attrs, inputs, out_ndims):
    """
    Shape function for take op.
    """
    if attrs.axis is None:
        return [_take_no_axis_shape_func(inputs[1], out_ndims[0])]
    axis = get_const_int(attrs.axis)
    batch_dims = get_const_int(attrs.batch_dims)
    data_ndim = int(inputs[0].shape[0])
    if inputs[1].shape:
        indices_ndim = int(inputs[1].shape[0])
    if axis < 0:
        axis += data_ndim
    assert 0 <= axis < data_ndim
    if batch_dims < 0:
        batch_dims += indices_ndim
    return [_take_with_axis_shape_func(*inputs, convert(axis), convert(batch_dims), out_ndims[0])]


@_reg.register_legalize("take")
def legalize_dyn_topk(attrs, inputs, types):
    """Legalize take op.
    Parameters
    ----------
    attrs : tvm.ir.Attrs
        Attributes of current op
    inputs : list of tvm.relay.Expr
        The args of the Relay expr to be legalized
    types : list of types
        List of input and output types
    Returns
    -------
    result : tvm.relay.Expr
        The legalized expr
    """
    return topi.take_legalize(attrs, inputs, types)


@script
def _argwhere_shape_func_1d(condition):
    out = output_tensor((2,), "int64")
    out[0] = int64(0)
    out[1] = int64(1)
    for i1 in range(condition.shape[0]):
        if condition[i1] != 0:
            out[0] += int64(1)
    return out


@script
def _argwhere_shape_func_2d(condition):
    out = output_tensor((2,), "int64")
    out[0] = int64(0)
    out[1] = int64(2)
    for i1 in range(condition.shape[0]):
        for i2 in range(condition.shape[1]):
            if condition[i1, i2] != 0:
                out[0] += int64(1)
    return out


@script
def _argwhere_shape_func_3d(condition):
    out = output_tensor((2,), "int64")
    out[0] = int64(0)
    out[1] = int64(3)
    for i1 in range(condition.shape[0]):
        for i2 in range(condition.shape[1]):
            for i3 in range(condition.shape[2]):
                if condition[i1, i2, i3] != 0:
                    out[0] += int64(1)
    return out


@script
def _argwhere_shape_func_4d(condition):
    out = output_tensor((2,), "int64")
    out[0] = int64(0)
    out[1] = int64(4)
    for i1 in range(condition.shape[0]):
        for i2 in range(condition.shape[1]):
            for i3 in range(condition.shape[2]):
                for i4 in range(condition.shape[3]):
                    if condition[i1, i2, i3, i4] != 0:
                        out[0] += int64(1)
    return out


@script
def _argwhere_shape_func_5d(condition):
    out = output_tensor((2,), "int64")
    out[0] = int64(0)
    out[1] = int64(5)
    for i1 in range(condition.shape[0]):
        for i2 in range(condition.shape[1]):
            for i3 in range(condition.shape[2]):
                for i4 in range(condition.shape[3]):
                    for i5 in range(condition.shape[4]):
                        if condition[i1, i2, i3, i4, i5] != 0:
                            out[0] += int64(1)
    return out


@_reg.register_shape_func("argwhere", True)
def argwhere_shape_func(attrs, inputs, out_ndims):
    """
    Shape function for argwhere.
    """
    if len(inputs[0].shape) == 1:
        return [_argwhere_shape_func_1d(inputs[0])]
    if len(inputs[0].shape) == 2:
        return [_argwhere_shape_func_2d(inputs[0])]
    if len(inputs[0].shape) == 3:
        return [_argwhere_shape_func_3d(inputs[0])]
    if len(inputs[0].shape) == 4:
        return [_argwhere_shape_func_4d(inputs[0])]
    if len(inputs[0].shape) == 5:
        return [_argwhere_shape_func_5d(inputs[0])]
    return ValueError("Does not support rank higher than 5 in argwhere")


_reg.register_shape_func("scatter", False, elemwise_shape_func)
_reg.register_shape_func("scatter_add", False, elemwise_shape_func)
_reg.register_shape_func("scatter_nd", False, elemwise_shape_func)


@script
def _sparse_fill_empty_rows_shape_func(sparse_indices, dense_shape):

    new_sparse_indices_shape = output_tensor((2,), "int64")
    new_sparse_values_shape = output_tensor((1,), "int64")
    empty_row_indicator_shape = output_tensor((1,), "int64")
    num_dense_rows = int64(dense_shape[0])

    if int64(sparse_indices.shape[0]) == int64(0):  # Handle Empty Case
        #  Total rows will equal dense_shape[0]
        new_sparse_indices_shape[0] = num_dense_rows
        new_sparse_indices_shape[1] = int64(sparse_indices.shape[1])
        new_sparse_values_shape[0] = num_dense_rows
        empty_row_indicator_shape[0] = num_dense_rows
        return (new_sparse_indices_shape, new_sparse_values_shape, empty_row_indicator_shape)

    else:
        count = int64(sparse_indices.shape[0])  # Add count of all rows already in sparse_indices
        for i in range(1, int64(sparse_indices.shape[0])):
            index = int64(sparse_indices[i, 0])
            prev_index = int64(sparse_indices[i - 1, 0] + 1)

            if index > prev_index:
                count += index - prev_index  # Add count of all rows between two consecutive indices

        count += int64(sparse_indices[0, 0])  # Add count from 0 to first row id in sparse_indices
        count += int64(
            num_dense_rows - 1 - sparse_indices[sparse_indices.shape[0] - 1, 0]
        )  # Add count from last row id to dense_shape - 1
        new_sparse_indices_shape[0] = int64(count)
        new_sparse_indices_shape[1] = int64(sparse_indices.shape[1])
        new_sparse_values_shape[0] = int64(count)
        empty_row_indicator_shape[0] = num_dense_rows
        return (new_sparse_indices_shape, new_sparse_values_shape, empty_row_indicator_shape)


@_reg.register_shape_func("sparse_fill_empty_rows", True)
def sparse_fill_empty_rows_func(attrs, inputs, _):
    return _sparse_fill_empty_rows_shape_func(inputs[0], inputs[2])


@script
def _sparse_reshape_shape_func(sparse_indices_shape, prev_shape_shape, new_shape_shape):
    indices_shape = output_tensor((2,), "int64")
    indices_shape[0] = int64(sparse_indices_shape[0])
    indices_shape[1] = int64(new_shape_shape[0])
    shape_tensor = output_tensor((1,), "int64")
    shape_tensor[0] = int64(new_shape_shape[0])
    return (indices_shape, shape_tensor)


@_reg.register_shape_func("sparse_reshape", False)
def sparse_reshape_shape_func(attrs, inputs, _):
    """
    Shape func for sparse_reshape.
    """
    return _sparse_reshape_shape_func(inputs[0], inputs[1], inputs[2])


@script
def _layout_transform_shape_func(
    data_shape, out_layout_len, dst_equal_list, dst_mul_list, dst_div_list, dst_mix_list
):
    out = output_tensor((out_layout_len,), "int64")
    for i in const_range(len(dst_equal_list)):
        out[dst_equal_list[i][0]] = data_shape[dst_equal_list[i][1]]
    for i in const_range(len(dst_mul_list)):
        out[dst_mul_list[i][0]] = data_shape[dst_mul_list[i][1]] * data_shape[dst_mul_list[i][2]]
    for i in const_range(len(dst_div_list)):
        out[dst_div_list[i][0]] = data_shape[dst_div_list[i][1]] // dst_div_list[i][3]
        out[dst_div_list[i][2]] = int64(dst_div_list[i][3])
    for i in const_range(len(dst_mix_list)):
        out[dst_mix_list[i][0]] = (
            data_shape[dst_mix_list[i][1]] * dst_mix_list[i][2] // dst_mix_list[i][4]
        )
        out[dst_mix_list[i][3]] = int64(dst_mix_list[i][4])

    return out


@_reg.register_shape_func("layout_transform", False)
def layout_transform_shape_func(attrs, inputs, _):
    """
    Shape function for layout_transform op.
    """

    def _fetch_axis(layout):
        major_axes = []
        minor_axes = {}
        num_start = -1
        for i, item in enumerate(layout):
            if "A" <= item <= "Z":
                major_axes.append(item)
            elif "a" <= item <= "z":
                last_num = int(layout[num_start:i])
                minor_axes[item] = last_num
                num_start = -1
            elif num_start < 0:
                num_start = i
        return major_axes, minor_axes

    _, src_minor_axes = _fetch_axis(attrs.src_layout)
    dst_major_axes, dst_minor_axes = _fetch_axis(attrs.dst_layout)
    src_letter_list = []
    dst_letter_list = []
    for item in attrs.src_layout:
        if "A" <= item <= "Z" or "a" <= item <= "z":
            src_letter_list.append(item)
    for item in attrs.dst_layout:
        if "A" <= item <= "Z" or "a" <= item <= "z":
            dst_letter_list.append(item)
    out_layout_len = len(dst_major_axes) + len(dst_minor_axes)
    dst_equal_list = []
    dst_mul_list = []
    dst_div_list = []
    dst_mix_list = []

    for key in dst_major_axes:
        if key.lower() not in dst_minor_axes:
            if key.lower() not in src_minor_axes:
                dst_equal_list.append((dst_letter_list.index(key), src_letter_list.index(key)))
            else:
                dst_mul_list.append(
                    (
                        dst_letter_list.index(key),
                        src_letter_list.index(key),
                        src_letter_list.index(key.lower()),
                    )
                )
        else:
            if key.lower() not in src_minor_axes:
                dst_div_list.append(
                    (
                        dst_letter_list.index(key),
                        src_letter_list.index(key),
                        dst_letter_list.index(key.lower()),
                        dst_minor_axes[key.lower()],
                    )
                )
            else:
                dst_mix_list.append(
                    (
                        dst_letter_list.index(key),
                        src_letter_list.index(key),
                        src_minor_axes[key.lower()],
                        dst_letter_list.index(key.lower()),
                        dst_minor_axes[key.lower()],
                    )
                )

    return [
        _layout_transform_shape_func(
            inputs[0],
            convert(out_layout_len),
            convert(dst_equal_list),
            convert(dst_mul_list),
            convert(dst_div_list),
            convert(dst_mix_list),
        )
    ]


@script
def _expand_dim_shape_func(data_shape, ndim, axis, num_newaxis):
    out = output_tensor((ndim + num_newaxis,), "int64")
    for i in const_range(out.shape[0]):
        if i < axis:
            out[i] = data_shape[i]
        elif i < axis + num_newaxis:
            out[i] = int64(1)
        else:
            out[i] = data_shape[i - num_newaxis]

    return out


@_reg.register_shape_func("expand_dims", False)
def expand_dim_shape_func(attrs, inputs, _):
    """
    Shape function for expand_dim op.
    """
    axis = get_const_int(attrs.axis)
    num_newaxis = get_const_int(attrs.num_newaxis)
    if axis < 0:
        axis = inputs[0].shape[0] + axis + 1
    ndim = inputs[0].shape[0] if inputs[0].shape else 0
    return [_expand_dim_shape_func(inputs[0], convert(ndim), convert(axis), convert(num_newaxis))]


@script
def _transpose_shape_func(data_shape, axes):
    out = output_tensor((data_shape.shape[0],), "int64")
    for i in const_range(len(axes)):
        out[i] = data_shape[axes[i]]

    return out


@_reg.register_shape_func("transpose", False)
def transpose_shape_func(attrs, inputs, _):
    """
    Shape function for transpose op.
    """
    axes = attrs.axes if attrs.axes is None else get_const_tuple(attrs.axes)
    if axes is None:
        axes = list(range(inputs[0].shape[0].value))
        axes.reverse()
    axes = list(axes)
    for i, axis in enumerate(axes):
        if axis < 0:
            axes[i] = inputs[0].shape[0] + axis
    return [_transpose_shape_func(inputs[0], convert(axes))]


_reg.register_schedule("transpose", strategy.schedule_transpose)


@script
def _squeeze_shape_func(data_shape, keep_axes, remove_axes):
    out = output_tensor((len(keep_axes),), "int64")
    for i in const_range(len(keep_axes)):
        out[i] = data_shape[keep_axes[i]]

    for i in const_range(len(remove_axes)):
        assert data_shape[remove_axes[i]] == 1, "Removed dimension must have size 1"

    return out


@_reg.register_shape_func("squeeze", False)
def squeeze_shape_func(attrs, inputs, _):
    """
    Shape function for squeeze op.
    """
    axis = attrs.axis if attrs.axis is None else get_const_tuple(attrs.axis)
    keep_axes = []
    remove_axes = []
    if axis is not None:
        for i in range(inputs[0].shape[0].value):
            if i not in axis:
                keep_axes.append(i)
            else:
                remove_axes.append(i)

    # Due to current relay type system, it is possible even
    # a static kernel function needs shape function. To handle
    # this case, we allow axis to be None in squeeze shape func
    # for now.
    # TODO(kevinthesun): Enhance relay type system to avoid this.
    if keep_axes:
        out = _squeeze_shape_func(inputs[0], convert(keep_axes), convert(remove_axes))
    else:
        out = te.compute((), lambda *indices: 0)
    return [out]


@script
def _reshape_like_shape_func(target_shape):
    out = output_tensor((target_shape.shape[0],), "int64")
    for i in const_range(target_shape.shape[0]):
        out[i] = target_shape[i]

    return out


@_reg.register_shape_func("reshape_like", False)
def reshape_like_shape_func(attrs, inputs, _):
    """
    Shape function for reshape_like op.
    """
    return [_reshape_like_shape_func(inputs[1])]


@script
def _tile_shape_func(data, reps, ndim, tndim, rndim):
    out = output_tensor((tndim,), "int64")

    if ndim == rndim:
        for i in const_range(tndim):
            out[i] = data[i] * int64(reps[i])
    elif ndim > rndim:
        ngap = ndim - rndim
        for i in const_range(ndim):
            if i < ngap:
                out[i] = data[i]
            else:
                out[i] = data[i] * int64(reps[i - ngap])
    else:
        rgap = rndim - ndim
        for i in const_range(rndim):
            if i < rgap:
                out[i] = int64(reps[i])
            else:
                out[i] = int64(reps[i]) * data[i - rgap]
    return out


@_reg.register_shape_func("tile", False)
def tile_shape_func(attrs, inputs, _):
    """
    Shape function for tile op.
    """
    reps = get_const_tuple(attrs.reps)
    ndim = inputs[0].shape[0].value
    rndim = len(reps)
    tndim = ndim if ndim > rndim else rndim
    return [
        _tile_shape_func(inputs[0], convert(reps), convert(ndim), convert(tndim), convert(rndim))
    ]


@script
def _split_shape_func(data_shape, index, indices_or_sections, param_is_indices, axis):
    out = output_tensor((data_shape.shape[0],), "int64")
    if param_is_indices:
        for i in const_range(data_shape.shape[0]):
            if i == axis:
                assert (
                    data_shape[axis] % indices_or_sections[0] == 0
                ), "num_sections must be an integer factor of the size of axis"
                out[i] = ceil_div(data_shape[axis], indices_or_sections[0])
            else:
                out[i] = data_shape[i]
    else:
        start = int64(0)
        if index > 0:
            start = int64(indices_or_sections[index - 1])
        end = data_shape[axis]
        if index < len(indices_or_sections):
            end = int64(indices_or_sections[index])
        for i in const_range(data_shape.shape[0]):
            if i == axis:
                out[i] = end - start
            else:
                out[i] = data_shape[i]
    return out


@_reg.register_shape_func("split", False)
def split_shape_func(attrs, inputs, _):
    """
    Shape function for split op.
    """
    if isinstance(attrs.indices_or_sections, (int, tvm.tir.IntImm)):
        indices_or_sections = get_const_int(attrs.indices_or_sections)
        assert indices_or_sections > 0, "Slice count must be > 0"
    else:
        indices_or_sections = list(get_const_tuple(attrs.indices_or_sections))
        assert sorted(indices_or_sections)[0] > 0 and indices_or_sections == sorted(
            indices_or_sections
        ), "split_indices must be sorted"

    axis = get_const_int(attrs.axis)

    if axis < 0:
        axis += get_const_int(inputs[0].shape[0])

    num_out = (
        indices_or_sections
        if isinstance(indices_or_sections, int)
        else len(indices_or_sections) + 1
    )

    param_is_indices = isinstance(indices_or_sections, int)
    if param_is_indices:
        indices_or_sections = [indices_or_sections]
    return [
        _split_shape_func(
            inputs[0],
            convert(i),
            convert(indices_or_sections),
            convert(param_is_indices),
            convert(axis),
        )
        for i in range(num_out)
    ]


@script
def _repeat_shape_func(data_shape, repeats, axis):
    out = output_tensor((data_shape.shape[0],), "int64")

    for i in const_range(data_shape.shape[0]):
        if i == axis:
            out[i] = int64(data_shape[i] * repeats)
        else:
            out[i] = data_shape[i]

    return out


@_reg.register_shape_func("repeat", False)
def repeat_shape_func(attrs, inputs, _):
    """
    Shape func for repeat.
    """
    axis = get_const_int(attrs.axis)
    if axis < 0:
        axis = inputs[0].shape[0] + axis
    return [_repeat_shape_func(inputs[0], attrs.repeats, convert(axis))]


@_reg.register_shape_func("broadcast_to_like", False)
def broadcast_to_like_shape_func(attrs, inputs, _):
    """
    Shape func for broadcast_to_like.
    """
    return [topi.math.identity(inputs[1])]


@script
def _stack_shape_func(data_shape, axis, num_inputs):
    out = output_tensor((data_shape.shape[0] + 1,), "int64")

    for i in const_range(data_shape.shape[0] + 1):
        if i == axis:
            out[i] = int64(num_inputs)
        elif i < axis:
            out[i] = data_shape[i]
        else:
            out[i] = data_shape[i - 1]

    return out


@_reg.register_shape_func("stack", False)
def stack_shape_func(attrs, inputs, _):
    """
    Shape func for stack.
    """
    axis = get_const_int(attrs.axis)
    if axis < 0:
        axis += inputs[0].shape[0] + 1
    return [_stack_shape_func(inputs[0], convert(axis), convert(len(inputs)))]


@script
def _broadcast_shape_tensors(shape_tensor1, shape_tensor2):
    rank1 = shape_tensor1.shape[0]
    rank2 = shape_tensor2.shape[0]
    out_rank = max(rank1, rank2)
    bcast_shape_tensor = output_tensor((out_rank,), "int64")

    for index in const_range(out_rank):
        dim1 = int64(1)
        dim2 = int64(1)

        if rank1 == out_rank:
            dim1 = shape_tensor1[index]
        elif rank1 - (out_rank - index) >= 0:
            dim1 = shape_tensor1[rank1 - (out_rank - index)]

        if rank2 == out_rank:
            dim2 = shape_tensor2[index]
        elif rank2 - (out_rank - index) >= 0:
            dim2 = shape_tensor2[rank2 - (out_rank - index)]

        assert dim1 == dim2 or dim1 == 1 or dim2 == 1, "Invalid broadcast shapes"
        bcast_shape_tensor[index] = max(dim1, dim2)

    return bcast_shape_tensor


@_reg.register_shape_func("where", False)
def where_shape_func(attrs, inputs, _):
    """
    Shape func for where.
    """

    def ensure_tensor(tensor):
        if len(tensor.shape) == 0:
            return topi.full((1,), "int64", 1)
        return tensor

    cond_shape = ensure_tensor(inputs[0])
    x_shape = ensure_tensor(inputs[1])
    y_shape = ensure_tensor(inputs[2])

    bcast_shape = _broadcast_shape_tensors(x_shape, y_shape)
    out_shape = _broadcast_shape_tensors(bcast_shape, cond_shape)

    return [out_shape]


@script
def _adv_index_post_process(data_shape, bcast_shape, num_indices):
    data_rank = data_shape.shape[0]
    bcast_rank = bcast_shape.shape[0]
    out = output_tensor((data_rank + bcast_rank - num_indices,), "int64")

    for i in const_range(bcast_rank):
        out[i] = bcast_shape[i]
    for i in const_range(data_rank - num_indices):
        out[i + bcast_rank] = data_shape[i + num_indices]
    return out


@_reg.register_shape_func("adv_index", False)
def adv_index_shape_func(attrs, inputs, _):
    """
    Shape func for adv_index.
    """
    bcast_shape = inputs[1]
    for i in inputs[2:]:
        bcast_shape = _broadcast_shape_tensors(bcast_shape, i)
    return [_adv_index_post_process(inputs[0], bcast_shape, convert(len(inputs) - 1))]


@script
def _unique_shape(data_shape):
    unique_shape = output_tensor((1,), "int64")
    indices_shape = output_tensor((1,), "int64")
    inverse_indices_shape = output_tensor((1,), "int64")
    num_unique_shape = output_tensor((1,), "int64")
    unique_shape[0] = data_shape[0]
    indices_shape[0] = data_shape[0]
    inverse_indices_shape[0] = data_shape[0]
    num_unique_shape[0] = int64(1)
    return (unique_shape, indices_shape, inverse_indices_shape, num_unique_shape)


@script
def _unique_with_counts_shape(data_shape):
    unique_shape = output_tensor((1,), "int64")
    indices_shape = output_tensor((1,), "int64")
    inverse_indices_shape = output_tensor((1,), "int64")
    num_unique_shape = output_tensor((1,), "int64")
    counts_shape = output_tensor((1,), "int64")
    unique_shape[0] = data_shape[0]
    indices_shape[0] = data_shape[0]
    inverse_indices_shape[0] = data_shape[0]
    num_unique_shape[0] = int64(1)
    counts_shape[0] = data_shape[0]
    return (unique_shape, indices_shape, inverse_indices_shape, num_unique_shape, counts_shape)


@_reg.register_shape_func("unique", False)
def unique_shape_func(attrs, inputs, _):
    """
    Shape func for unique operator.
    """
    if attrs.return_counts:
        return _unique_with_counts_shape(inputs[0])
    else:
        return _unique_shape(inputs[0])


@script
def _gather_nd_shape(data_shape, indices_shape, batch_dims, index_rank):
    ndim = data_shape.shape[0]
    # using mdim = indices_shape[0] wouldn't work because a rank cannot
    # depend on a runtime shape dimension of indices tensor, even if the
    # dimension is always a known, fixed value. As a workaround, we assume that
    # the fixed gather dimension (the size of an indexing tuple) is recorded
    # in gather_nd op attributes.
    mdim = index_rank
    kdim = indices_shape.shape[0] - 1
    out_shape = output_tensor((kdim + ndim - (mdim + batch_dims),), "int64")
    for i in range(1, kdim + 1):
        out_shape[i - 1] = indices_shape[i]
    for i in range(mdim + batch_dims, ndim):
        out_shape[kdim + i - (mdim + batch_dims)] = data_shape[i]
    return out_shape


@_reg.register_shape_func("gather_nd", False)
def gather_nd_shape_func(attrs, inputs, _):
    """
    Shape func for gather_nd operator.
    """
    batch_dims = get_const_int(attrs.batch_dims)
    index_rank = get_const_int(attrs.index_rank)

    assert index_rank > 0, "index_rank needs to be specified for dynamic gather_nd"

    return [_gather_nd_shape(inputs[0], inputs[1], convert(batch_dims), convert(index_rank))]


@script
def _gather_shape(data_shape, indices_shape, axis):
    out_shape = output_tensor((data_shape.shape[0],), "int64")
    for i in range(data_shape.shape[0]):
        if i != axis:
            assert (
                data_shape[i] == indices_shape[i]
            ), "data and indices size at non-gather axes must be the same"
        out_shape[i] = indices_shape[i]
    return out_shape


@_reg.register_shape_func("gather", False)
def gather_shape_func(attrs, inputs, _):
    """
    Shape func for gather operator.
    """
    return [_gather_shape(inputs[0], inputs[1], attrs.axis)]