/*
 * 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.
 */

/*!
 * \brief NN op constructions
 * \file topi/nn.h
 */
#ifndef TVM_TOPI_NN_H_
#define TVM_TOPI_NN_H_

#include <tvm/arith/analyzer.h>
#include <tvm/te/operation.h>
#include <tvm/tirx/expr.h>
#include <tvm/tirx/op.h>
#include <tvm/topi/detail/constant_utils.h>
#include <tvm/topi/reduction.h>
#include <tvm/topi/tags.h>
#include <tvm/topi/transform.h>

#include <algorithm>
#include <string>

namespace tvm {
namespace topi {

using namespace tvm::te;

/*!
 * \brief Creates an operation that performs a rectified linear unit
 *
 * \param t The input tensor
 * \param threshold The relu threshold (default 0)
 * \param name The name of the operation
 * \param tag The tag to mark the operation
 *
 * \return A Tensor whose op member is the relu operation
 */
template <typename T>
inline tvm::te::Tensor relu(const tvm::te::Tensor& t, T threshold = static_cast<T>(0),
                            std::string name = "T_relu", std::string tag = kElementWise) {
  return tvm::te::compute(
      t->shape,
      [&](const tvm::ffi::Array<tvm::tirx::PrimVar>& i) {
        auto threshold_const = tvm::tirx::MakeConst(tvm::PrimType(t->dtype), threshold);
        return tvm::max(t(i), threshold_const);
      },
      name, tag);
}

/*!
 * \brief Creates an operation that performs a leaky rectified linear unit
 *
 * \param t The input tensor
 * \param alpha The slope for the small gradient when t < 0
 * \param name The name of the operation
 * \param tag The tag to mark the operation
 *
 * \return A Tensor whose op member is the leaky relu operation
 */
inline tvm::te::Tensor leaky_relu(const tvm::te::Tensor& t, double alpha = 0.1,
                                  std::string name = "T_leaky_relu",
                                  std::string tag = kElementWise) {
  return tvm::te::compute(
      t->shape,
      [&](const tvm::ffi::Array<tvm::tirx::PrimVar>& i) {
        auto value = t(i);
        auto calpha = tvm::tirx::MakeConst(value.ty(), alpha);
        return tvm::tirx::Select(value > 0, value, value * calpha);
      },
      name, tag);
}

/*!
 * \brief Creates an operation that performs a parametric rectified linear unit
 *
 * \param x The input data tensor
 * \param slope The channel-wise slope tensor
 * \param axis The axis where the channel data needs to be applied
 * \param name The name of the operation
 * \param tag The tag to mark the operation
 *
 * \return A Tensor whose op member is the parametric relu operation
 */
inline tvm::te::Tensor prelu(const tvm::te::Tensor& x, const tvm::te::Tensor& slope,
                             const int axis = 1, std::string name = "T_prelu",
                             std::string tag = kBroadcast) {
  TVM_FFI_ICHECK((size_t)axis < x->shape.size()) << "Wrong axis (" << axis << ")value. ";
  TVM_FFI_ICHECK(topi::detail::GetConstInt(slope->shape[0]) ==
                 topi::detail::GetConstInt(x->shape[axis]))
      << "Wrong slope shape received.";

  return tvm::te::compute(
      x->shape,
      [&](const tvm::ffi::Array<tvm::tirx::PrimVar>& indices) {
        auto xval = x(indices);
        return tvm::tirx::Select(xval > 0, xval, xval * slope(indices[axis]));
      },
      name, tag);
}

/*!
 * \brief Creates an operation that performs padding
 *
 * \param t The input tensor
 * \param pad_before An Array of Expr describing the padding before the
 * respective iterator
 * \param pad_after An Array of Expr describing the padding after the
 * respective iterator
 * \param pad_value The value to fill padding elements with
 * \param pad_mode Padding type to use.
 * "constant" pads with constant_value;
 * "edge" pads using the edge values of the input array;
 * "reflect" pads by reflecting values with respect to the edges.
 * \param dyn_output_shape Output shape of the pad op, default nullptr.
 * You only need to pass this in if the shape was evaluated dynamically.
 * \param name The name of the operation
 * \param tag The tag to mark the operation
 *
 * \return A Tensor whose op member is the padding operation
 *
 * \note
 *  The pad_after Array must either be empty or have the same length as
 *  pad_before
 *  When pad_after is empty, it takes the same values as pad_before (symmetric
 *  padding)
 *  The pad Array applies from the leading dimensions and skips missing
 *  trailing dimensions:
 *
 *      pad(t(i, j, k), {1}, {0}) returns the equivalent operation for
 *          the following pseudocode:
 *              for i in [1, t.shape[0] + 2]:
 *                  for i in [1, t.shape[0] + 2]:
 *                      for i in [1, t.shape[0] + 2]:
 *                         name(i,j,k) =
 *                             (1 <= i <= t.shape[0] + 1) ?
 *                                 t(i-1, j, k) : 0;
 *
 *
 */
inline tvm::te::Tensor pad(
    const tvm::te::Tensor& t, const tvm::ffi::Array<tvm::PrimExpr>& pad_before,
    tvm::ffi::Array<tvm::PrimExpr> pad_after = tvm::ffi::Array<tvm::PrimExpr>(),
    PrimExpr pad_value = PrimExpr(), std::string name = "T_pad", std::string tag = kElementWise,
    std::string pad_mode = "constant", const ffi::Array<PrimExpr>* dyn_output_shape = nullptr) {
  if (pad_after.size() < pad_before.size()) {
    for (size_t i = pad_after.size(); i < pad_before.size(); ++i) {
      pad_after.push_back(pad_before[i]);
    }
  }

  arith::Analyzer analyzer;
  TVM_FFI_ICHECK_GE(pad_before.size(), 1);
  TVM_FFI_ICHECK_EQ(pad_before.size(), pad_after.size());
  tvm::ffi::Array<tvm::PrimExpr> pad_before_int32;
  tvm::ffi::Array<tvm::PrimExpr> pad_after_int32;

  for (const auto& ele : pad_before) {
    pad_before_int32.push_back(tvm::cast(tvm::PrimType::Int(32), ele));
  }
  for (const auto& ele : pad_after) {
    pad_after_int32.push_back(tvm::cast(tvm::PrimType::Int(32), ele));
  }

  tvm::ffi::Array<tvm::PrimExpr> output_shape;
  if (dyn_output_shape == nullptr) {
    for (size_t i = 0; i < t->shape.size(); ++i) {
      if (i >= pad_before.size()) {
        output_shape.push_back(t->shape[i]);
      } else {
        output_shape.push_back(
            analyzer->Simplify(t->shape[i] + pad_before_int32[i] + pad_after_int32[i]));
      }
    }
  } else {
    for (size_t i = 0; i < dyn_output_shape->size(); i++) {
      output_shape.push_back((*dyn_output_shape)[i]);
    }
  }

  if (!pad_value.defined()) {
    pad_value = tvm::tirx::MakeConst(tvm::PrimType(t->dtype), 0);
  }

  auto l = [&](tvm::ffi::Array<tvm::tirx::PrimVar> ovars) {
    tvm::ffi::Array<tvm::PrimExpr> indices;
    tvm::ffi::Array<tvm::PrimExpr> sel;
    tvm::ffi::Array<tvm::PrimExpr> pad_idx;
    for (size_t i = 0; i < t->shape.size(); ++i) {
      if (i >= pad_before_int32.size()) {
        indices.push_back(ovars[i]);
        continue;
      }
      if (!topi::detail::EqualCheck(pad_before_int32[i], 0)) {
        sel.push_back(ovars[i] >= pad_before_int32[i]);
        indices.push_back(ovars[i] - pad_before_int32[i]);
      } else {
        indices.push_back(ovars[i]);
      }
      if (!topi::detail::EqualCheck(pad_after_int32[i], 0)) {
        sel.push_back(analyzer->Simplify(ovars[i].as_or_throw<PrimExpr>() <
                                         pad_before_int32[i] + t->shape[i]));
      }
      if (pad_mode == "edge") {
        pad_idx.push_back(
            tvm::if_then_else(ovars[i].as_or_throw<PrimExpr>() < pad_before[i], 0,
                              tvm::if_then_else(ovars[i] >= pad_before[i] + t->shape[i],
                                                t->shape[i] - 1, ovars[i] - pad_before[i])));
      } else if (pad_mode == "reflect") {
        pad_idx.push_back(tvm::if_then_else(
            ovars[i].as_or_throw<PrimExpr>() < pad_before[i], pad_before[i] - ovars[i],
            tvm::if_then_else(ovars[i] >= pad_before[i] + t->shape[i],
                              t->shape[i] * 2 - ovars[i] + pad_before[i] - 2,
                              ovars[i] - pad_before[i])));
      }
    }
    if (sel.size() != 0) {
      if (pad_mode == "constant") {
        return tvm::if_then_else(
            foldl([](PrimExpr a, PrimExpr b, Span span) { return tvm::logical_and(a, b, span); },
                  IntImm::Bool(true), sel),
            t(indices), pad_value);
      } else if (pad_mode == "edge" || pad_mode == "reflect") {
        return tvm::if_then_else(
            foldl([](PrimExpr a, PrimExpr b, Span span) { return tvm::logical_and(a, b, span); },
                  IntImm::Bool(true), sel),
            t(indices), t(pad_idx));
      }
    }
    return t(indices);
  };
  return tvm::te::compute(output_shape, l, name, tag);
}

/*!
 * \brief Creates an operation that performs a 2-D convolution with an
 * NCHW-layout
 *
 * \param I The 4-D input tensor
 * \param W The 4-D weight tensor
 * \param pad_h A static constant padding amount applied to the height of the
 * image, before and after (symmetric padding)
 * \param pad_w A static constant padding amount applied to the width of the
 * image, before and after (symmetric padding)
 * \param stride_h A static constant striding amount applied to the height of
 * the image, before and after (symmetric padding)
 * \param stride_w A static constant strindingamount applied to the width of
 * the image, before and after (symmetric padding)
 * \param name The name of the operation
 * \param tag The tag to mark the operation
 *
 * \return A Tensor whose op member is the 2-D convolution operation (NCHW
 * layout)
 */
inline tvm::te::Tensor conv2d_nchw(const tvm::te::Tensor& I, const tvm::te::Tensor& W,
                                   int pad_h = 0, int pad_w = 0, int stride_h = 1, int stride_w = 1,
                                   std::string name = "T_conv2d_nchw",
                                   std::string tag = kConv2dNCHW) {
  TVM_FFI_ICHECK_EQ(4, I->shape.size());
  TVM_FFI_ICHECK_EQ(4, W->shape.size());
  auto pH = I->shape[2];
  auto pW = I->shape[3];
  tvm::ffi::Array<tvm::PrimExpr> output_shape{
      I->shape[0],                                                    // B
      W->shape[0],                                                    // O
      indexdiv(I->shape[2] - W->shape[2] + 2 * pad_h, stride_h) + 1,  // H
      indexdiv(I->shape[3] - W->shape[3] + 2 * pad_w, stride_w) + 1   // W
  };
  auto i = tvm::te::reduce_axis(tvm::Range{0, I->shape[1]}, "i");
  auto kh = tvm::te::reduce_axis(tvm::Range{0, W->shape[2]}, "kh");
  auto kw = tvm::te::reduce_axis(tvm::Range{0, W->shape[3]}, "kw");
  auto T =
      (pad_h == 0 && pad_w == 0) ? I : pad(I, {tvm::PrimExpr(0), tvm::PrimExpr(0), pad_h, pad_w});
  auto l = [&](tvm::tirx::PrimVar b, tvm::tirx::PrimVar o, tvm::tirx::PrimVar h,
               tvm::tirx::PrimVar w) {
    return tvm::sum(T(b, i, stride_h * h + kh, stride_w * w + kw) * W(o, i, kh, kw), {i, kh, kw});
  };
  return tvm::te::compute(output_shape, l, name, tag);
}

/*!
 * \brief Creates an operation for 2-D convolution layer with an HWCN-layout
 *
 * \param I The 4-D input tensor
 * \param W The 4-D weight tensor
 * \param pad_h A static constant padding amount applied to the height of the
 * image, before and after (symmetric padding)
 * \param pad_w A static constant padding amount applied to the width of the
 * image, before and after (symmetric padding)
 * \param stride_h A static constant striding amount applied to the height of
 * the image, before and after (symmetric padding)
 * \param stride_w A static constant strindingamount applied to the width of
 * the image, before and after (symmetric padding)
 * \param name The name of the operation
 * \param tag The tag to mark the operation
 *
 * \return A Tensor whose op member is the 2-D convolution operation
 * (HWCN layout)
 */
inline tvm::te::Tensor conv2d_hwcn(const tvm::te::Tensor& I, const tvm::te::Tensor& W,
                                   int pad_h = 0, int pad_w = 0, int stride_h = 1, int stride_w = 1,
                                   std::string name = "T_conv2d_hwcn",
                                   std::string tag = kConv2dHWCN) {
  TVM_FFI_ICHECK_EQ(4, I->shape.size());
  TVM_FFI_ICHECK_EQ(4, W->shape.size());
  auto pH = I->shape[2];
  auto pW = I->shape[3];
  tvm::ffi::Array<tvm::PrimExpr> output_shape{
      indexdiv(I->shape[2] - W->shape[2] + 2 * pad_h, stride_h) + 1,  // H
      indexdiv(I->shape[3] - W->shape[3] + 2 * pad_w, stride_w) + 1,  // W
      I->shape[2],                                                    // B
      W->shape[3]                                                     // O
  };
  auto i = tvm::te::reduce_axis(tvm::Range{0, I->shape[3]}, "i");
  auto kh = tvm::te::reduce_axis(tvm::Range{0, W->shape[0]}, "kh");
  auto kw = tvm::te::reduce_axis(tvm::Range{0, W->shape[1]}, "kw");
  auto T = (pad_h == 0 && pad_w == 0) ? I : pad(I, {pad_h, pad_w});
  auto l = [&](tvm::tirx::PrimVar b, tvm::tirx::PrimVar o, tvm::tirx::PrimVar h,
               tvm::tirx::PrimVar w) {
    return tvm::sum(T(stride_h * h + kh, stride_w * w + kw, i, b) * W(kh, kw, i, o), {i, kh, kw});
  };
  return tvm::te::compute(output_shape, l, name, tag);
}

/*!
 * \brief Creates an operation that performs a 2-D depthwise convolution with
 * an NCHW-layout
 *
 * \param I The 4-D input tensor
 * \param W The 4-D weight tensor
 * \param pad_h A static constant padding amount applied to the height of the
 * image, before and after (symmetric padding)
 * \param pad_w A static constant padding amount applied to the width of the
 * image, before and after (symmetric padding)
 * \param stride_h A static constant striding amount applied to the height of
 * the image, before and after (symmetric padding)
 * \param stride_w A static constant strindingamount applied to the width of
 * the image, before and after (symmetric padding)
 * \param name The name of the operation
 * \param tag The tag to mark the operation
 *
 * \return A Tensor whose op member is the 2-D depthwise convolution operation
 * (NCHW layout)
 */
inline tvm::te::Tensor depthwise_conv2d_nchw(const tvm::te::Tensor& I, const tvm::te::Tensor& W,
                                             int pad_h = 0, int pad_w = 0, int stride_h = 1,
                                             int stride_w = 1,
                                             std::string name = "T_depthwise_conv2d_nchw",
                                             std::string tag = kDepthwiseConv2dNCHW) {
  TVM_FFI_ICHECK_EQ(4, I->shape.size());
  TVM_FFI_ICHECK_EQ(4, W->shape.size());
  auto pH = I->shape[2];
  auto pW = I->shape[3];
  auto pCM = W->shape[1];  // channel_multiplier
  tvm::ffi::Array<tvm::PrimExpr> output_shape{
      I->shape[0],                                                    // B
      W->shape[1],                                                    // O
      indexdiv(I->shape[2] - W->shape[2] + 2 * pad_h, stride_h) + 1,  // H
      indexdiv(I->shape[3] - W->shape[3] + 2 * pad_w, stride_w) + 1   // W
  };
  auto i = tvm::te::reduce_axis(tvm::Range{0, I->shape[1]}, "i");
  auto kh = tvm::te::reduce_axis(tvm::Range{0, W->shape[2]}, "kh");
  auto kw = tvm::te::reduce_axis(tvm::Range{0, W->shape[3]}, "kw");
  auto T =
      (pad_h == 0 && pad_w == 0) ? I : pad(I, {tvm::PrimExpr(0), tvm::PrimExpr(0), pad_h, pad_w});
  auto l = [&](tvm::tirx::PrimVar b, tvm::tirx::PrimVar o, tvm::tirx::PrimVar h,
               tvm::tirx::PrimVar w) {
    return tvm::sum(T(b, indexdiv(i, pCM), stride_h * h + kh, stride_w * w + kw) *
                        W(indexdiv(i, pCM), indexmod(o, pCM), kh, kw),
                    {i, kh, kw});
  };
  return tvm::te::compute(output_shape, l, name, tag);
}

inline tvm::te::Tensor depthwise_conv2d_nhwc(const tvm::te::Tensor& I, const tvm::te::Tensor& W,
                                             int pad_h = 0, int pad_w = 0, int stride_h = 1,
                                             int stride_w = 1,
                                             std::string name = "T_depthwise_conv2d_nhwc",
                                             std::string tag = kDepthwiseConv2dNHWC) {
  TVM_FFI_ICHECK_EQ(4, I->shape.size());
  TVM_FFI_ICHECK_EQ(4, W->shape.size());
  auto pH = I->shape[1];
  auto pW = I->shape[2];
  auto pCM = W->shape[1];  // channel_multiplier
  tvm::ffi::Array<tvm::PrimExpr> output_shape{
      I->shape[0],                                                    // B
      indexdiv(I->shape[1] - W->shape[1] + 2 * pad_h, stride_h) + 1,  // H
      indexdiv(I->shape[2] - W->shape[2] + 2 * pad_w, stride_w) + 1,  // W
      W->shape[3],                                                    // O
  };
  auto i = tvm::te::reduce_axis(tvm::Range{0, I->shape[3]}, "i");
  auto kh = tvm::te::reduce_axis(tvm::Range{0, W->shape[0]}, "kh");
  auto kw = tvm::te::reduce_axis(tvm::Range{0, W->shape[1]}, "kw");
  auto T =
      (pad_h == 0 && pad_w == 0) ? I : pad(I, {tvm::PrimExpr(0), pad_h, pad_w, tvm::PrimExpr(0)});
  auto l = [&](tvm::tirx::PrimVar b, tvm::tirx::PrimVar h, tvm::tirx::PrimVar w,
               tvm::tirx::PrimVar o) {
    return tvm::sum(T(b, stride_h * h + kh, stride_w * w + kw, indexdiv(i, pCM)) *
                        W(kh, kw, indexdiv(i, pCM), indexmod(o, pCM)),
                    {kh, kw, i});
  };
  return tvm::te::compute(output_shape, l, name, tag);
}

/*!
 * \brief Creates an operation that performs a 2-D group convolution with
 * an NGCHW-layout
 *
 * \param I The 5-D input tensor
 * \param W The 5-D weight tensor
 * \param pad_h A static constant padding amount applied to the height of the
 * image, before and after (symmetric padding)
 * \param pad_w A static constant padding amount applied to the width of the
 * image, before and after (symmetric padding)
 * \param stride_h A static constant striding amount applied to the height of
 * the image, before and after (symmetric padding)
 * \param stride_w A static constant strindingamount applied to the width of
 * the image, before and after (symmetric padding)
 * \param name The name of the operation
 * \param tag The tag to mark the operation
 *
 * \return A Tensor whose op member is the 2-D groupconvolution operation
 * (NCHW layout)
 */
inline tvm::te::Tensor group_conv2d_ngchw(const tvm::te::Tensor& I, const tvm::te::Tensor& W,
                                          int pad_h = 0, int pad_w = 0, int stride_h = 1,
                                          int stride_w = 1,
                                          std::string name = "T_group_conv2d_ngchw",
                                          std::string tag = kGroupConv2d) {
  TVM_FFI_ICHECK_EQ(5, I->shape.size());
  TVM_FFI_ICHECK_EQ(5, W->shape.size());
  auto pH = I->shape[2];
  auto pW = I->shape[3];
  tvm::ffi::Array<tvm::PrimExpr> output_shape{
      I->shape[0],                                                    // B
      I->shape[1],                                                    // G
      W->shape[2],                                                    // O
      indexdiv(I->shape[3] - W->shape[3] + 2 * pad_h, stride_h) + 1,  // H
      indexdiv(I->shape[4] - W->shape[4] + 2 * pad_w, stride_w) + 1   // W
  };
  auto i = tvm::te::reduce_axis(tvm::Range{0, I->shape[2]}, "i");
  auto kh = tvm::te::reduce_axis(tvm::Range{0, W->shape[3]}, "kh");
  auto kw = tvm::te::reduce_axis(tvm::Range{0, W->shape[4]}, "kw");

  auto T = (pad_h == 0 && pad_w == 0)
               ? I
               : pad(I, {tvm::PrimExpr(0), tvm::PrimExpr(0), tvm::PrimExpr(0), pad_h, pad_w});
  auto l = [&](tvm::ffi::Array<tvm::tirx::PrimVar> args) {
    tvm::tirx::PrimVar b = args[0];
    tvm::tirx::PrimVar g = args[1];
    tvm::tirx::PrimVar o = args[2];
    tvm::tirx::PrimVar h = args[3];
    tvm::tirx::PrimVar w = args[4];
    return tvm::sum(I(b, g, i, stride_h * h + kh, stride_w * w + kw) * W(g, i, o, kh, kw),
                    {i, kh, kw});
  };
  return tvm::te::compute(output_shape, l, name, tag);
}

/*!
 * \brief Divide spatial dimensions of the input into a grid of blocks.
 *
 * \param data The input tensor.
 * \param block_shape The size of the spatial block.
 * \param pad_before The zero-padding size before each spatial dimension.
 * \param pad_after The zero-padding size after each spatial dimension.
 * \param pad_value The value used for padding.
 * \param name The name of the operation.
 * \param tag The tag to mark the operation.
 *
 * \return A Tensor whose op member is the space_to_batch_nd operation
 */
inline tvm::te::Tensor space_to_batch_nd(const tvm::te::Tensor& data,
                                         const tvm::ffi::Array<int64_t>& block_shape,
                                         const tvm::ffi::Array<tvm::PrimExpr>& pad_before,
                                         const tvm::ffi::Array<tvm::PrimExpr>& pad_after,
                                         PrimExpr pad_value = PrimExpr(),
                                         std::string name = "space_to_batch_nd",
                                         std::string tag = kInjective) {
  tvm::te::Tensor padded_t;
  TVM_FFI_ICHECK_EQ(pad_before.size(), pad_after.size());
  TVM_FFI_ICHECK_EQ(block_shape.size(), pad_before.size())
      << "Paddings must be provided for each spatial dimension";
  tvm::ffi::Array<tvm::PrimExpr> pad_before_int32;
  tvm::ffi::Array<tvm::PrimExpr> pad_after_int32;

  // pad size for batch dimension is 0
  pad_before_int32.push_back(tvm::cast(tvm::PrimType::Int(32), 0));
  pad_after_int32.push_back(tvm::cast(tvm::PrimType::Int(32), 0));
  // insert pad sizes given for spatial dimensions
  for (const auto& ele : pad_before) {
    pad_before_int32.push_back(tvm::cast(tvm::PrimType::Int(32), ele));
  }
  for (const auto& ele : pad_after) {
    pad_after_int32.push_back(tvm::cast(tvm::PrimType::Int(32), ele));
  }

  // pad the input with paddings provided
  if (!pad_value.defined()) {
    pad_value = tvm::tirx::MakeConst(tvm::PrimType(data->dtype), 0);
  }
  padded_t = pad(data, pad_before_int32, pad_after_int32, pad_value);

  auto input_shape = data->shape;
  auto padded_shape = padded_t->shape;

  // infer shapes
  tvm::ffi::Array<PrimExpr> r_shape;
  tvm::ffi::Array<int64_t> axis;
  tvm::ffi::Array<PrimExpr> o_shape;

  size_t num_block_dims = block_shape.size();
  int batch = static_cast<int>(GetConstInt(input_shape[0]));
  tvm::PrimExpr block_shape_prod(1);
  r_shape.push_back(batch);

  for (size_t i = 1; i <= num_block_dims; i++) {
    int padded_input = static_cast<int>(GetConstInt(padded_shape[i]));
    int block_size = static_cast<int>(block_shape[i - 1]);
    TVM_FFI_ICHECK_EQ((padded_input % block_size), 0)
        << "(" << i
        << ")th "
           "Input dimension after padding ("
        << padded_input << ")"
        << " must be divisible by its block size (" << block_size << ")";

    PrimExpr bs = IntImm::Int64(block_shape[i - 1]);
    r_shape.push_back(div(padded_shape[i], bs));
    r_shape.push_back(bs);
    block_shape_prod *= bs;
    axis.push_back(static_cast<int64_t>(r_shape.size() - 1));  // index of block_shape[i - 1]
  }

  size_t n = axis.size();
  axis.push_back(0);  // batch is at index 0
  // index of (padded_shape[i] / block_shape[i - 1]) in r_shape
  for (size_t i = 0; i < n; i++) {
    axis.push_back(axis[i] - 1);
  }
  o_shape.push_back(tvm::PrimExpr(batch) * block_shape_prod);
  for (size_t i = 1; i <= num_block_dims; i++) {
    PrimExpr bs = IntImm::Int64(block_shape[i - 1]);
    o_shape.push_back(div(padded_shape[i], bs));
  }
  // append remaining shape
  for (size_t i = num_block_dims + 1; i < input_shape.size(); i++) {
    r_shape.push_back(input_shape[i]);
    axis.push_back(
        static_cast<int64_t>(r_shape.size() - 1));  // index of remaining shape in r_shape
    o_shape.push_back(input_shape[i]);
  }

  tvm::te::Tensor output = reshape(padded_t, r_shape);
  output = transpose(output, axis);
  output = reshape(output, o_shape);

  return output;
}

/*!
 * \brief Reshape the batch dimension into spatial dimensions.
 *
 * \param data The input tensor.
 * \param block_shape The size of the spatial block.
 * \param crop_begin_list The begin crop size for each spatial dimension.
 * \param crop_end_list The end crop size for each spatial dimension.
 * \param name The name of the operation.
 * \param tag The tag to mark the operation.
 *
 * \return A Tensor whose op member is the batch_to_space_nd operation
 */
inline tvm::te::Tensor batch_to_space_nd(const tvm::te::Tensor& data,
                                         const tvm::ffi::Array<int64_t>& block_shape,
                                         const tvm::ffi::Array<tvm::PrimExpr>& crop_begin_list,
                                         const tvm::ffi::Array<tvm::PrimExpr>& crop_end_list,
                                         std::string name = "batch_to_space_nd",
                                         std::string tag = kInjective) {
  // Construct shapes for reshape and transpose operation
  ffi::Array<PrimExpr> in_shape = data->shape;
  ffi::Array<PrimExpr> r_shape;
  ffi::Array<int64_t> axis;
  size_t num_block_dims = block_shape.size();
  size_t num_input_dims = in_shape.size();
  tvm::PrimExpr block_shape_prod(1);
  int batch = static_cast<int>(GetConstInt(in_shape[0]));

  for (size_t i = 0; i < num_block_dims; i++) {
    PrimExpr bs = IntImm::Int64(block_shape[i]);
    r_shape.push_back(bs);
    block_shape_prod *= bs;
  }
  axis.push_back(static_cast<int64_t>(r_shape.size()));  // axis of (batch / block_shape_prod)
  r_shape.push_back(batch / block_shape_prod);

  for (size_t i = 1; i < num_input_dims; i++) {
    axis.push_back(static_cast<int64_t>(r_shape.size()));  // axis of in_shape[i]
    if (axis.size() < (num_block_dims + num_input_dims)) {
      axis.push_back(
          static_cast<int64_t>(r_shape.size() - (num_block_dims + 1)));  // axis of block_shape[i]
    }
    r_shape.push_back(in_shape[i]);
  }

  ffi::Array<PrimExpr> r_p_shape;
  r_p_shape.push_back(batch / block_shape_prod);
  for (size_t i = 1; i <= num_block_dims; i++) {
    PrimExpr bs = IntImm::Int64(block_shape[i - 1]);
    r_p_shape.push_back(in_shape[i] * bs);
  }
  for (size_t i = num_block_dims + 1; i < num_input_dims; i++) {
    r_p_shape.push_back(in_shape[i]);
  }

  tvm::te::Tensor out;
  out = reshape(data, r_shape);
  out = transpose(out, axis);
  out = reshape(out, r_p_shape);

  // Crop the start and end of dimensions of out
  ffi::Array<ffi::Optional<IntImm>> begin_idx, end_idx;
  ffi::Array<IntImm> strides;
  PrimType index_ty = PrimType::Int(64);
  for (size_t i = 0; i < r_p_shape.size(); ++i) {
    strides.push_back(IntImm(index_ty, 1));
    if (i > 0 && i <= num_block_dims) {
      // prepare begin and end index for spatial dimensions
      int64_t begin_i = GetConstInt(crop_begin_list[i - 1]);
      int64_t end_i = GetConstInt(crop_end_list[i - 1]);
      int64_t out_i = GetConstInt(r_p_shape[i]);
      TVM_FFI_ICHECK_GT(out_i, (begin_i + end_i))
          << "Incorrect crop sizes for (" << i << ")th dim, can not crop more than"
          << " output size" << out_i << " vs " << (begin_i + end_i);
      begin_idx.push_back(IntImm(index_ty, begin_i));
      end_idx.push_back(IntImm(index_ty, out_i - end_i));
    } else {
      // ignore the batch and remaining dimension
      begin_idx.push_back(IntImm(index_ty, 0));
      end_idx.push_back(IntImm(index_ty, GetConstInt(r_p_shape[i])));
    }
  }

  out = strided_slice(out, begin_idx, end_idx, strides);
  return out;
}

/*!
 * \brief Negative log likelihood loss.
 *
 * \param predictions The prediction tensor.
 * \param targets The target tensor.
 * \param weights A manual rescaling weight given to each class.
 * \param reduction The reduction method to apply to the output.
 * \param ignore_index The target value to ignore.
 * \param name The name of the operation.
 * \param tag The tag to mark the operation.
 *
 * \return The negative log likelihood loss of the predictions and targets.
 */
inline Tensor nll_loss(const Tensor& predictions, const Tensor& targets, const Tensor& weights,
                       std::string reduction = "mean", int ignore_index = -100,
                       const std::string name = "nll_loss", const std::string tag = kBroadcast) {
  if (predictions.ndim() == 1) {
    // corner case: no batch in shape
    // prediction->shape = (C,), targets->shape = (), weights->shape = (C,)
    auto T = tvm::te::compute(
        {},
        [&](const tvm::ffi::Array<tvm::tirx::PrimVar>& target_indices) {
          auto c = targets();
          return tvm::tirx::Select(c != ignore_index, -predictions(c) * weights(c),
                                   tvm::tirx::MakeConst(tvm::PrimType(predictions->dtype), 0));
        },
        name, tag);
    if (reduction == "mean") {
      auto W = tvm::te::compute(
          {},
          [&](const tvm::ffi::Array<tvm::tirx::PrimVar>& target_indices) {
            auto c = targets();
            return tvm::tirx::Select(c != ignore_index, weights(c),
                                     tvm::tirx::MakeConst(tvm::PrimType(predictions->dtype), 0));
          },
          name, tag);
      return topi::divide(T, W);
    } else {
      return T;
    }
  }
  auto T = tvm::te::compute(
      targets->shape,
      [&](const tvm::ffi::Array<tvm::tirx::PrimVar>& target_indices) {
        auto c = targets(target_indices);
        tvm::ffi::Array<tvm::PrimExpr> pred_indices;
        pred_indices.push_back(target_indices[0]);  // batch index
        pred_indices.push_back(c);                  // class index
        for (size_t i = 1; i < target_indices.size(); i++) {
          pred_indices.push_back(target_indices[i]);  // indices for multidimensional loss
        }
        return tvm::tirx::Select(c != ignore_index, -predictions(pred_indices) * weights(c),
                                 tvm::tirx::MakeConst(tvm::PrimType(predictions->dtype), 0));
      },
      name, tag);
  TVM_FFI_ICHECK(T->shape.size() != 0);
  if (reduction == "mean") {
    auto W = tvm::te::compute(
        targets->shape,
        [&](const tvm::ffi::Array<tvm::tirx::PrimVar>& target_indices) {
          auto c = targets(target_indices);
          return tvm::tirx::Select(c != ignore_index, weights(c),
                                   tvm::tirx::MakeConst(tvm::PrimType(predictions->dtype), 0));
        },
        name, tag);
    return topi::divide(topi::sum(T, tvm::ffi::Array<int64_t>(nullptr)),
                        topi::sum(W, tvm::ffi::Array<int64_t>(nullptr)));
  } else if (reduction == "sum") {
    return topi::sum(T, tvm::ffi::Array<int64_t>(nullptr));
  } else {  // reduction == "none"
    return T;
  }
}

}  // namespace topi
}  // namespace tvm
#endif  // TVM_TOPI_NN_H_