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

/*!
 * \file tvm/relax/attrs/nn.h
 * \brief Attributes for neural network operators.
 */
#ifndef TVM_RELAX_ATTRS_NN_H_
#define TVM_RELAX_ATTRS_NN_H_

#include <tvm/relax/expr.h>

namespace tvm {
namespace relax {

/*! \brief Attributes used in Conv1d operator */
struct Conv1DAttrs : public AttrsNode {
  ffi::Array<int64_t> strides;
  ffi::Array<int64_t> padding;
  ffi::Array<int64_t> dilation;
  int groups;
  ffi::String data_layout;
  ffi::String kernel_layout;
  ffi::String out_layout;
  ffi::Optional<DLDataType> out_dtype;

  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<Conv1DAttrs>()
        .def_ro("strides", &Conv1DAttrs::strides, "Specifies the strides of the convolution.")
        .def_ro("padding", &Conv1DAttrs::padding,
                "If padding is non-zero, then the input is implicitly zero-padded"
                "Padding support both symmetric and asymmetric as"
                "one int : same padding used on both sides"
                "two int : padding width in the order of (left, right)")
        .def_ro("dilation", &Conv1DAttrs::dilation,
                "Specifies the dilation rate to use for dilated convolution.")
        .def_ro("groups", &Conv1DAttrs::groups,
                "Number of groups to split the input into for grouped convolution. The number of "
                "input and "
                "output channels should be divisible by the number of groups.")
        .def_ro("data_layout", &Conv1DAttrs::data_layout,
                "Dimension ordering of input data. Can be 'NCW', 'NWC', etc."
                "'N', 'C', 'W' stands for batch, channel, width"
                "dimensions respectively. Convolution is applied on the 'W' dimensions.")
        .def_ro("kernel_layout", &Conv1DAttrs::kernel_layout,
                "Dimension ordering of weight. Can be 'OIW', 'IOW', etc."
                "'O', 'I', 'W' stands for num_filter, input_channel, and width"
                "dimensions respectively.")
        .def_ro("out_layout", &Conv1DAttrs::out_layout,
                "Dimension ordering of output. Can be 'NCW', 'NWC', etc."
                "'N', 'C', 'W' stands for batch, channel, and width"
                "dimensions respectively. Default to be same as input layout.")
        .def_ro("out_dtype", &Conv1DAttrs::out_dtype,
                "Output data type, set to explicit type under mixed precision setting");
  }
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.Conv1DAttrs", Conv1DAttrs, AttrsNode);
};  // struct Conv1dAttrs

/*! \brief Attributes used in Conv2d operator */
struct Conv2DAttrs : public AttrsNode {
  ffi::Array<int64_t> strides;
  ffi::Array<int64_t> padding;
  ffi::Array<int64_t> dilation;
  int groups;
  ffi::String data_layout;
  ffi::String kernel_layout;
  ffi::String out_layout;
  ffi::Optional<DLDataType> out_dtype;

  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<Conv2DAttrs>()
        .def_ro("strides", &Conv2DAttrs::strides, "Specifies the strides of the convolution.")
        .def_ro("padding", &Conv2DAttrs::padding,
                "If padding is non-zero, then the input is implicitly zero-padded"
                "Padding support both symmetric and asymmetric as"
                "one int : same padding used on all sides"
                "two int : bottom, right will use same padding as top, left"
                "four int : padding width in the order of (top, left, bottom, right)")
        .def_ro("dilation", &Conv2DAttrs::dilation,
                "Specifies the dilation rate to use for dilated convolution.")
        .def_ro("groups", &Conv2DAttrs::groups,
                "Number of groups to split the input into for grouped convolution. The number of "
                "input and "
                "output channels should be divisible by the number of groups.")
        .def_ro("data_layout", &Conv2DAttrs::data_layout,
                "Dimension ordering of input data. Can be 'NCHW', 'NHWC', etc."
                "'N', 'C', 'H', 'W' stands for batch, channel, height, and width"
                "dimensions respectively. Convolution is applied on the 'H' and"
                "'W' dimensions.")
        .def_ro("kernel_layout", &Conv2DAttrs::kernel_layout,
                "Dimension ordering of weight. Can be 'OIHW', 'OIHW16o16i', etc."
                "'O', 'I', 'H', 'W' stands for num_filter, input_channel, height, and width"
                "dimensions respectively.")
        .def_ro("out_layout", &Conv2DAttrs::out_layout,
                "Dimension ordering of output. Can be 'NCHW', 'NHWC', etc."
                "'N', 'C', 'H', 'W' stands for batch, channel, height, and width"
                "dimensions respectively. Default to be same as input layout.")
        .def_ro("out_dtype", &Conv2DAttrs::out_dtype,
                "Output data type, set to explicit type under mixed precision setting");
  }
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.Conv2DAttrs", Conv2DAttrs, AttrsNode);
};  // struct Conv2dAttrs

/*! \brief Attributes used in Conv3d operator */
struct Conv3DAttrs : public AttrsNode {
  ffi::Array<int64_t> strides;
  ffi::Array<int64_t> padding;
  ffi::Array<int64_t> dilation;
  int groups;
  ffi::String data_layout;
  ffi::String kernel_layout;
  ffi::String out_layout;
  ffi::Optional<DLDataType> out_dtype;

  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<Conv3DAttrs>()
        .def_ro("strides", &Conv3DAttrs::strides, "Specifies the strides of the convolution.")
        .def_ro(
            "padding", &Conv3DAttrs::padding,
            "If padding is non-zero, then the input is implicitly zero-padded"
            "Padding support both symmetric and asymmetric as"
            "one int : same padding used on all sides"
            "two int : bottom, right will use same padding as top, left"
            "four int : padding width in the order of (forward, back, top, left, bottom, right)")
        .def_ro("dilation", &Conv3DAttrs::dilation,
                "Specifies the dilation rate to use for dilated convolution.")
        .def_ro("groups", &Conv3DAttrs::groups,
                "Number of groups to split the input into for grouped convolution. The number of "
                "input and "
                "output channels should be divisible by the number of groups.")
        .def_ro("data_layout", &Conv3DAttrs::data_layout,
                "Dimension ordering of input data. Can be 'NCDHW', 'NDHWC', etc."
                "'N', 'C', 'D', 'H', 'W' stands for batch, channel, depth, height, and width"
                "dimensions respectively. Convolution is applied on the 'D', 'H', and"
                "'W' dimensions.")
        .def_ro(
            "kernel_layout", &Conv3DAttrs::kernel_layout,
            "Dimension ordering of weight. Can be 'OIDHW', 'OIDHW16o16i', etc."
            "'O', 'I', 'D', 'H', 'W' stands for num_filter, input_channel, depth, height, and width"
            "dimensions respectively.")
        .def_ro("out_layout", &Conv3DAttrs::out_layout,
                "Dimension ordering of output. Can be 'NCDHW', 'NDHWC', etc."
                "'N', 'C', 'D', 'H', 'W' stands for batch, channel, depth, height, and width"
                "dimensions respectively. Default to be same as input layout.")
        .def_ro("out_dtype", &Conv3DAttrs::out_dtype,
                "Output data type, set to explicit type under mixed precision setting");
  }
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.Conv3DAttrs", Conv3DAttrs, AttrsNode);
};  // struct Conv3dAttrs

/*! \brief Attributes used in Conv1DTranspose operator */
struct Conv1DTransposeAttrs : public AttrsNode {
  ffi::Array<int64_t> strides;
  ffi::Array<int64_t> padding;
  ffi::Array<int64_t> output_padding;
  ffi::Array<int64_t> dilation;
  int groups;
  ffi::String data_layout;
  ffi::String kernel_layout;
  ffi::String out_layout;
  ffi::Optional<DLDataType> out_dtype;

  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<Conv1DTransposeAttrs>()
        .def_ro("strides", &Conv1DTransposeAttrs::strides,
                "Specifies the strides of the convolution.")
        .def_ro("padding", &Conv1DTransposeAttrs::padding,
                "If padding is non-zero, then the input is implicitly zero-padded"
                "Padding support both symmetric and asymmetric as"
                "one int : same padding used on both sides"
                "two int : padding width in the order of (left, right)")
        .def_ro("output_padding", &Conv1DTransposeAttrs::output_padding,
                "Used to disambiguate the output shape.")
        .def_ro("dilation", &Conv1DTransposeAttrs::dilation,
                "Specifies the dilation rate to use for dilated convolution.")
        .def_ro("groups", &Conv1DTransposeAttrs::groups,
                "Number of groups to split the input into for grouped convolution. The number of "
                "input and "
                "output channels should be divisible by the number of groups.")
        .def_ro("data_layout", &Conv1DTransposeAttrs::data_layout,
                "Dimension ordering of input data. Can be 'NCW', 'NWC', etc."
                "'N', 'C', 'W' stands for batch, channel, width"
                "dimensions respectively. Convolution is applied on the 'W' dimensions.")
        .def_ro("kernel_layout", &Conv1DTransposeAttrs::kernel_layout,
                "Dimension ordering of weight. Can be 'OIW', 'IOW', etc."
                "'O', 'I', 'W' stands for num_filter, input_channel, and width"
                "dimensions respectively.")
        .def_ro("out_layout", &Conv1DTransposeAttrs::out_layout,
                "Dimension ordering of output. Can be 'NCW', 'NWC', etc."
                "'N', 'C', 'W' stands for batch, channel, and width"
                "dimensions respectively. Default to be same as input layout.")
        .def_ro("out_dtype", &Conv1DTransposeAttrs::out_dtype,
                "Output data type, set to explicit type under mixed precision setting");
  }
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.Conv1DTransposeAttrs", Conv1DTransposeAttrs,
                                    AttrsNode);
};  // struct Conv1DTransposeAttrs

/*! \brief Attributes used in Conv2d operator */
struct Conv2DTransposeAttrs : public AttrsNode {
  ffi::Array<int64_t> strides;
  ffi::Array<int64_t> padding;
  ffi::Array<int64_t> output_padding;
  ffi::Array<int64_t> dilation;
  int groups;
  ffi::String data_layout;
  ffi::String kernel_layout;
  ffi::String out_layout;
  ffi::Optional<DLDataType> out_dtype;

  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<Conv2DTransposeAttrs>()
        .def_ro("strides", &Conv2DTransposeAttrs::strides,
                "Specifies the strides of the convolution.")
        .def_ro("padding", &Conv2DTransposeAttrs::padding,
                "If padding is non-zero, then the input is implicitly zero-padded"
                "Padding support both symmetric and asymmetric as"
                "one int : same padding used on all sides"
                "two int : bottom, right will use same padding as top, left"
                "four int : padding width in the order of (top, left, bottom, right)")
        .def_ro("output_padding", &Conv2DTransposeAttrs::output_padding,
                "Used to disambiguate the output shape.")
        .def_ro("dilation", &Conv2DTransposeAttrs::dilation,
                "Specifies the dilation rate to use for dilated convolution.")
        .def_ro("groups", &Conv2DTransposeAttrs::groups,
                "Number of groups to split the input into for grouped convolution. The number of "
                "input and "
                "output channels should be divisible by the number of groups.")
        .def_ro("data_layout", &Conv2DTransposeAttrs::data_layout,
                "Dimension ordering of input data. Can be 'NCHW', 'NHWC', etc."
                "'N', 'C', 'H', 'W' stands for batch, channel, height, and width"
                "dimensions respectively. Convolution is applied on the 'H' and"
                "'W' dimensions.")
        .def_ro("kernel_layout", &Conv2DTransposeAttrs::kernel_layout,
                "Dimension ordering of weight. Can be 'OIHW', 'OIHW16o16i', etc."
                "'O', 'I', 'H', 'W' stands for num_filter, input_channel, height, and width"
                "dimensions respectively.")
        .def_ro("out_layout", &Conv2DTransposeAttrs::out_layout,
                "Dimension ordering of output. Can be 'NCHW', 'NHWC', etc."
                "'N', 'C', 'H', 'W' stands for batch, channel, height, and width"
                "dimensions respectively. Default to be same as input layout.")
        .def_ro("out_dtype", &Conv2DTransposeAttrs::out_dtype,
                "Output data type, set to explicit type under mixed precision setting");
  }
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.Conv2DTransposeAttrs", Conv2DTransposeAttrs,
                                    AttrsNode);
};  // struct Conv2DTransposeAttrs

/*! \brief Attributes used in Conv3dTranspose operator */
struct Conv3DTransposeAttrs : public AttrsNode {
  ffi::Array<int64_t> strides;
  ffi::Array<int64_t> padding;
  ffi::Array<int64_t> output_padding;
  ffi::Array<int64_t> dilation;
  int groups;
  ffi::String data_layout;
  ffi::String kernel_layout;
  ffi::String out_layout;
  ffi::Optional<DLDataType> out_dtype;

  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<Conv3DTransposeAttrs>()
        .def_ro("strides", &Conv3DTransposeAttrs::strides,
                "Specifies the strides of the convolution.")
        .def_ro("padding", &Conv3DTransposeAttrs::padding,
                "If padding is non-zero, then the input is implicitly zero-padded"
                "Padding support both symmetric and asymmetric as"
                "one int : same padding used on all sides"
                "three int : back/bottom/right will use same padding as front/top/left"
                "six int : padding width in the order of (front, top, left, back, bottom, right)")
        .def_ro("output_padding", &Conv3DTransposeAttrs::output_padding,
                "Used to disambiguate the output shape.")
        .def_ro("dilation", &Conv3DTransposeAttrs::dilation,
                "Specifies the dilation rate to use for dilated convolution.")
        .def_ro("groups", &Conv3DTransposeAttrs::groups,
                "Number of groups to split the input into for grouped convolution. The number of "
                "input and "
                "output channels should be divisible by the number of groups.")
        .def_ro("data_layout", &Conv3DTransposeAttrs::data_layout,
                "Dimension ordering of input data. Can be 'NCDHW', 'NDHWC', etc."
                "'N', 'C', 'D', 'H', 'W' stands for batch, channel, depth, height, and width"
                "dimensions respectively. Convolution is applied on the 'D', 'H', and"
                "'W' dimensions.")
        .def_ro(
            "kernel_layout", &Conv3DTransposeAttrs::kernel_layout,
            "Dimension ordering of weight. Can be 'IODHW', etc."
            "'I', 'O', 'D', 'H', 'W' stands for input_channel, output_channel, depth, height, and "
            "width"
            "dimensions respectively.")
        .def_ro("out_layout", &Conv3DTransposeAttrs::out_layout,
                "Dimension ordering of output. Can be 'NCDHW', 'NDHWC', etc."
                "'N', 'C', 'D', 'H', 'W' stands for batch, channel, depth, height, and width"
                "dimensions respectively. Default to be same as input layout.")
        .def_ro("out_dtype", &Conv3DTransposeAttrs::out_dtype,
                "Output data type, set to explicit type under mixed precision setting");
  }
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.Conv3DTransposeAttrs", Conv3DTransposeAttrs,
                                    AttrsNode);
};  // struct Conv3DTransposeAttrs

/*! \brief Attributes used in max_pool1d and avg_pool1d operator */
struct Pool1DAttrs : public AttrsNode {
  ffi::Array<int64_t> pool_size;
  ffi::Array<int64_t> strides;
  ffi::Array<int64_t> padding;
  ffi::Array<int64_t> dilation;
  bool ceil_mode;
  bool count_include_pad;
  ffi::String layout;
  ffi::String out_layout;

  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<Pool1DAttrs>()
        .def_ro("pool_size", &Pool1DAttrs::pool_size, "Size of the pooling windows.")
        .def_ro("strides", &Pool1DAttrs::strides, "Specifies the strides of the convolution.")
        .def_ro("dilation", &Pool1DAttrs::dilation, "Specifies the dilation of the convolution.")
        .def_ro("padding", &Pool1DAttrs::padding,
                "If padding is non-zero, then the input is implicitly zero-padded"
                "Padding support both symmetric and asymmetric as"
                "one int : same padding used on all sides"
                "two int : padding width in the order of (left, right)")
        .def_ro(
            "ceil_mode", &Pool1DAttrs::ceil_mode,
            "A boolean indicating if use ceil or floor to compute the output shape. By using ceil, "
            "every element in the input tensor will be covered by a sliding window.")
        .def_ro("count_include_pad", &Pool1DAttrs::count_include_pad,
                "When true, will include padding to compute the average")
        .def_ro("layout", &Pool1DAttrs::layout,
                "Dimension ordering of input data. Can be 'NCW', 'NWC', etc."
                "'N', 'C', 'W' stands for batch, channel, and width"
                "dimensions respectively. Pooling is applied on the 'W' dimensions.",
                refl::DefaultValue("NCW"))
        .def_ro("out_layout", &Pool1DAttrs::out_layout,
                "Dimension ordering of output data. Can be 'NCW', 'NWC', etc."
                "'N', 'C', 'W' stands for batch, channel, and width"
                "dimensions respectively. Pooling is applied on the 'W' dimensions.");
  }
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.Pool1DAttrs", Pool1DAttrs, AttrsNode);
};  // struct Pool1dAttrs

/*! \brief Attributes used in max_pool2d and avg_pool2d operator */
struct Pool2DAttrs : public AttrsNode {
  ffi::Array<int64_t> pool_size;
  ffi::Array<int64_t> strides;
  ffi::Array<int64_t> padding;
  ffi::Array<int64_t> dilation;
  bool ceil_mode;
  bool count_include_pad;
  ffi::String layout;
  ffi::String out_layout;

  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<Pool2DAttrs>()
        .def_ro("pool_size", &Pool2DAttrs::pool_size, "Size of the pooling windows.")
        .def_ro("strides", &Pool2DAttrs::strides, "Specifies the strides of the convolution.")
        .def_ro("dilation", &Pool2DAttrs::dilation, "Specifies the dilation of the convolution.")
        .def_ro("padding", &Pool2DAttrs::padding,
                "If padding is non-zero, then the input is implicitly zero-padded"
                "Padding support both symmetric and asymmetric as"
                "one int : same padding used on all sides"
                "two int : bottom, right will use same padding as top, left"
                "four int : padding width in the order of (top, left, bottom, right)")
        .def_ro(
            "ceil_mode", &Pool2DAttrs::ceil_mode,
            "A boolean indicating if use ceil or floor to compute the output shape. By using ceil, "
            "every element in the input tensor will be covered by a sliding window.")
        .def_ro("count_include_pad", &Pool2DAttrs::count_include_pad,
                "When true, will include padding to compute the average")
        .def_ro("layout", &Pool2DAttrs::layout,
                "Dimension ordering of input data. Can be 'NCHW', 'NHWC', etc."
                "'N', 'C', 'H', 'W' stands for batch, channel, height, and width"
                "dimensions respectively. Pooling is applied on the 'H' and"
                "'W' dimensions.")
        .def_ro("out_layout", &Pool2DAttrs::out_layout,
                "Dimension ordering of output data. Can be 'NCHW', 'NHWC', etc."
                "'N', 'C', 'H', 'W' stands for batch, channel, height, and width"
                "dimensions respectively. Pooling is applied on the 'H' and"
                "'W' dimensions.");
  }
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.Pool2DAttrs", Pool2DAttrs, AttrsNode);
};  // struct Pool2dAttrs

/*! \brief Attributes used in max_pool3d and avg_pool3d operator */
struct Pool3DAttrs : public AttrsNode {
  ffi::Array<int64_t> pool_size;
  ffi::Array<int64_t> strides;
  ffi::Array<int64_t> padding;
  ffi::Array<int64_t> dilation;
  bool ceil_mode;
  bool count_include_pad;
  ffi::String layout;
  ffi::String out_layout;

  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<Pool3DAttrs>()
        .def_ro("pool_size", &Pool3DAttrs::pool_size, "Size of the pooling windows.")
        .def_ro("strides", &Pool3DAttrs::strides, "Specifies the strides of the convolution.")
        .def_ro("dilation", &Pool3DAttrs::dilation, "Specifies the dilation of the convolution.")
        .def_ro("padding", &Pool3DAttrs::padding,
                "If padding is non-zero, then the input is implicitly zero-padded"
                "Padding support both symmetric and asymmetric as"
                "one int : same padding used on all sides"
                "three int : back, bottom, right will use same padding as front, top, left"
                "four int : padding width in the order of (front, top, left, back, bottom, right)")
        .def_ro(
            "ceil_mode", &Pool3DAttrs::ceil_mode,
            "A boolean indicating if use ceil or floor to compute the output shape. By using ceil, "
            "every element in the input tensor will be covered by a sliding window.")
        .def_ro("count_include_pad", &Pool3DAttrs::count_include_pad,
                "When true, will include padding to compute the average")
        .def_ro("layout", &Pool3DAttrs::layout,
                "Dimension ordering of input data. Can be 'NCDHW', 'NDHWC', etc."
                "'N', 'C', 'D', 'H', 'W' stands for batch, channel, depth, height, and width"
                "dimensions respectively. Pooling is applied on the 'D', 'H' and"
                "'W' dimensions.")
        .def_ro("out_layout", &Pool3DAttrs::out_layout,
                "Dimension ordering of output data. Can be 'NCDHW', 'NDHWC', etc."
                "'N', 'C', 'D', 'H', 'W' stands for batch, channel, depth, height, and width"
                "dimensions respectively. Pooling is applied on the 'D', 'H' and"
                "'W' dimensions.");
  }
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.Pool3DAttrs", Pool3DAttrs, AttrsNode);
};  // struct Pool3dAttrs

/*! \brief Attributes for 1d adaptive pool operator */
struct AdaptivePool1DAttrs : public AttrsNode {
  ffi::Optional<ffi::Array<int64_t>> output_size;
  ffi::String layout;
  ffi::String out_layout;

  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<AdaptivePool1DAttrs>()
        .def_ro("output_size", &AdaptivePool1DAttrs::output_size, "Output width.")
        .def_ro("layout", &AdaptivePool1DAttrs::layout,
                "Dimension ordering of input data. Can be 'NCW', 'NWC', etc."
                "'N', 'C', 'W' stands for batch, channel and width"
                "dimensions respectively. Pooling is applied on the"
                "'W' dimensions.")
        .def_ro("out_layout", &AdaptivePool1DAttrs::out_layout,
                "Dimension ordering of output data. Can be 'NCW', 'NWC', etc."
                "'N', 'C', 'W' stands for batch, channel and width"
                "dimensions respectively. Pooling is applied on the"
                "'W' dimensions.");
  }
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.AdaptivePool1DAttrs", AdaptivePool1DAttrs,
                                    AttrsNode);
};  // struct AdaptivePool1DAttrs

/*! \brief Attributes for 2d adaptive pool operator */
struct AdaptivePool2DAttrs : public AttrsNode {
  ffi::Optional<ffi::Array<int64_t>> output_size;
  ffi::String layout;
  ffi::String out_layout;

  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<AdaptivePool2DAttrs>()
        .def_ro("output_size", &AdaptivePool2DAttrs::output_size, "Output height and width.")
        .def_ro("layout", &AdaptivePool2DAttrs::layout,
                "Dimension ordering of input data. Can be 'NCHW', 'NHWC', etc."
                "'N', 'C', 'H', 'W' stands for batch, channel, height, and width"
                "dimensions respectively. Pooling is applied on the 'H' and"
                "'W' dimensions.")
        .def_ro("out_layout", &AdaptivePool2DAttrs::out_layout,
                "Dimension ordering of output data. Can be 'NCHW', 'NHWC', etc."
                "'N', 'C', 'H', 'W' stands for batch, channel, height, and width"
                "dimensions respectively. Pooling is applied on the 'H' and"
                "'W' dimensions.");
  }
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.AdaptivePool2DAttrs", AdaptivePool2DAttrs,
                                    AttrsNode);
};  // struct AdaptivePool2DAttrs

/*! \brief Attributes for 3d adaptive pool operator */
struct AdaptivePool3DAttrs : public AttrsNode {
  ffi::Optional<ffi::Array<int64_t>> output_size;
  ffi::String layout;
  ffi::String out_layout;

  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<AdaptivePool3DAttrs>()
        .def_ro("output_size", &AdaptivePool3DAttrs::output_size, "Output depth, height and width.")
        .def_ro("layout", &AdaptivePool3DAttrs::layout,
                "Dimension ordering of input data. Can be 'NCDHW', 'NDHWC', etc."
                "'N', 'C', 'D', 'H', 'W' stands for batch, channel, depth, height, and width"
                "dimensions respectively. Pooling is applied on 'D', 'H' and"
                "'W' dimensions.")
        .def_ro("out_layout", &AdaptivePool3DAttrs::out_layout,
                "Dimension ordering of output data. Can be 'NCDHW', 'NDHWC', etc."
                "'N', 'C', 'D', 'H', 'W' stands for batch, channel, depth, height, and width"
                "dimensions respectively. Pooling is applied on 'D', 'H' and"
                "'W' dimensions.");
  }
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.AdaptivePool3DAttrs", AdaptivePool3DAttrs,
                                    AttrsNode);
};  // struct AdaptivePool3DAttrs

/*! \brief Attributes used in softmax operators */
struct SoftmaxAttrs : public AttrsNode {
  int axis;

  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<SoftmaxAttrs>().def_ro("axis", &SoftmaxAttrs::axis,
                                           "The axis to sum over when computing softmax.");
  }
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.SoftmaxAttrs", SoftmaxAttrs, AttrsNode);
};

/*! \brief Attributes used in softmax operators */
struct LeakyReluAttrs : public AttrsNode {
  double alpha;

  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<LeakyReluAttrs>().def_ro("alpha", &LeakyReluAttrs::alpha,
                                             "The slope of the negative part.");
  }
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.LeakyReluAttrs", LeakyReluAttrs, AttrsNode);
};

/*! \brief Attributes used in softplus operators */
struct SoftplusAttrs : public AttrsNode {
  double beta;
  double threshold;

  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<SoftplusAttrs>()
        .def_ro("beta", &SoftplusAttrs::beta,
                "Scaling factor controlling the sharpness of the Softplus transition.")
        .def_ro("threshold", &SoftplusAttrs::threshold,
                "Value determining when to use linear approximation for numerical stability.");
  }
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.SoftplusAttrs", SoftplusAttrs, AttrsNode);
};

/*! \brief Attributes used in PReLU operator */
struct PReluAttrs : public AttrsNode {
  int axis;

  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<PReluAttrs>().def_ro("axis", &PReluAttrs::axis,
                                         "The axis along which the alpha values are applied.");
  }
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.PReluAttrs", PReluAttrs, AttrsNode);
};

/*! \brief Attributes used in batch_norm operator */
struct BatchNormAttrs : public AttrsNode {
  int axis;
  double epsilon;
  bool center;
  bool scale;
  double momentum;
  bool training;

  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<BatchNormAttrs>()
        .def_ro("axis", &BatchNormAttrs::axis, "The axis along which the normalization is applied.")
        .def_ro("epsilon", &BatchNormAttrs::epsilon,
                "Small float added to variance to avoid dividing by zero")
        .def_ro("center", &BatchNormAttrs::center,
                "Indicating if the beta offset will be added to the normalized tensor.")
        .def_ro("scale", &BatchNormAttrs::scale,
                "Indicating if the gamma scale will be multiplied.")
        .def_ro("momentum", &BatchNormAttrs::momentum,
                "The value used for the moving_mean and moving_var update.")
        .def_ro("training", &BatchNormAttrs::training,
                "Whether we are training (i.e., not in eval mode).");
  }
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.BatchNormAttrs", BatchNormAttrs, AttrsNode);
};  // struct BatchNormAttrs

/*! \brief Attributes used in layer_norm operator */
struct LayerNormAttrs : public AttrsNode {
  ffi::Array<int64_t> axes;
  double epsilon;
  bool center;
  bool scale;

  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<LayerNormAttrs>()
        .def_ro("axes", &LayerNormAttrs::axes,
                "The axes that along which the normalization is applied.")
        .def_ro("epsilon", &LayerNormAttrs::epsilon,
                "Small float added to variance to avoid dividing by zero")
        .def_ro("center", &LayerNormAttrs::center,
                "Indicating if the beta offset will be added to the normalized tensor.")
        .def_ro("scale", &LayerNormAttrs::scale,
                "Indicating if the gamma scale will be multiplied.");
  }
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.LayerNormAttrs", LayerNormAttrs, AttrsNode);
};  // struct LayerNormAttrs

/*! \brief Attributes used in group_norm operator */
struct GroupNormAttrs : public AttrsNode {
  int num_groups;
  int channel_axis;
  ffi::Array<int64_t> axes;
  double epsilon;
  bool center;
  bool scale;

  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<GroupNormAttrs>()
        .def_ro("num_groups", &GroupNormAttrs::num_groups,
                "The number of groups to separate the channels into.")
        .def_ro("channel_axis", &GroupNormAttrs::channel_axis,
                "The axis that represents the channel.")
        .def_ro(
            "axes", &GroupNormAttrs::axes,
            "The axes that along which the normalization is applied (excluding the channel axis).")
        .def_ro("epsilon", &GroupNormAttrs::epsilon,
                "Small float added to variance to avoid dividing by zero")
        .def_ro("center", &GroupNormAttrs::center,
                "Indicating if the beta offset will be added to the normalized tensor.")
        .def_ro("scale", &GroupNormAttrs::scale,
                "Indicating if the gamma scale will be multiplied.");
  }
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.GroupNormAttrs", GroupNormAttrs, AttrsNode);
};  // struct GroupNormAttrs

/*! \brief Attributes used in instance_norm operator */
struct InstanceNormAttrs : public AttrsNode {
  int channel_axis;
  ffi::Array<int64_t> axes;
  double epsilon;
  bool center;
  bool scale;

  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<InstanceNormAttrs>()
        .def_ro("channel_axis", &InstanceNormAttrs::channel_axis,
                "The axis that represents the channel.")
        .def_ro("axes", &InstanceNormAttrs::axes,
                "The axes that along which the normalization is applied.")
        .def_ro("epsilon", &InstanceNormAttrs::epsilon,
                "Small float added to variance to avoid dividing by zero")
        .def_ro("center", &InstanceNormAttrs::center,
                "Indicating if the beta offset will be added to the normalized tensor.")
        .def_ro("scale", &InstanceNormAttrs::scale,
                "Indicating if the gamma scale will be multiplied.");
  }
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.InstanceNormAttrs", InstanceNormAttrs, AttrsNode);
};  // struct InstanceNormAttrs

/*! \brief Attributes used in rms_norm operator */
struct RMSNormAttrs : public AttrsNode {
  ffi::Array<int64_t> axes;
  double epsilon;

  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<RMSNormAttrs>()
        .def_ro("axes", &RMSNormAttrs::axes,
                "The axes that along which the normalization is applied.")
        .def_ro("epsilon", &RMSNormAttrs::epsilon,
                "Small float added to variance to avoid dividing by zero");
  }
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.RMSNormAttrs", RMSNormAttrs, AttrsNode);
};  // struct RMSNormAttrs

/*! \brief Attributes used in nll_loss operator */
struct NLLLossAttrs : public AttrsNode {
  ffi::String reduction;
  int ignore_index;

  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<NLLLossAttrs>()
        .def_ro("reduction", &NLLLossAttrs::reduction,
                "The reduction method to apply to the output. Can be"
                "'none', 'mean' or 'sum'.",
                refl::DefaultValue("mean"))
        .def_ro("ignore_index", &NLLLossAttrs::ignore_index, "The target value to ignore.");
  }
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.NLLLossAttrs", NLLLossAttrs, AttrsNode);
};  // struct NLLLossAttrs

/*! \brief Attributes used in dropout operator */
struct DropoutAttrs : public AttrsNode {
  double rate;

  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<DropoutAttrs>().def_ro(
        "rate", &DropoutAttrs::rate,
        "Fraction of the input that gets dropped out during training time");
  }
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.DropoutAttrs", DropoutAttrs, AttrsNode);
};  // struct DropoutAttrs

/*! \brief Attributes used in Attention operator */
struct AttentionAttrs : public AttrsNode {
  ffi::Optional<FloatImm> scale;
  ffi::Optional<ffi::String> causal_mask;
  ffi::Optional<IntImm> window_size;

  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<AttentionAttrs>()
        .def_ro(
            "scale", &AttentionAttrs::scale,
            "The custom scale applied before the softmax. The default value is 1 / sqrt(head_dim).")
        .def_ro("causal_mask", &AttentionAttrs::causal_mask,
                "The type of the causal mask, i.e. 'TopLeft' and 'BottomRight'.")
        .def_ro("window_size", &AttentionAttrs::window_size,
                "The size of the window for sliding-window attention.");
  }
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.AttentionAttrs", AttentionAttrs, AttrsNode);
};  // struct AttentionAttrs

/*! \brief Attributes used for the padding operator */
struct PadAttrs : public AttrsNode {
  ffi::Array<int64_t> pad_width;
  double pad_value = 0.0;
  tvm::ffi::String pad_mode;

  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<PadAttrs>()
        .def_ro("pad_width", &PadAttrs::pad_width,
                "Number of values padded to the edges of each axis, "
                "in the format of (before_1, after_1, ..., before_N, after_N)")
        .def_ro("pad_value", &PadAttrs::pad_value, "The value to fill in padded area with",
                refl::DefaultValue(0.0))
        .def_ro("pad_mode", &PadAttrs::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.",
                refl::DefaultValue("constant"));
  }
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.PadAttrs", PadAttrs, AttrsNode);
};

/*! \brief Attributes used for the pixel shuffle operator */
struct PixelShuffleAttrs : public AttrsNode {
  int upscale_factor;

  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<PixelShuffleAttrs>().def_ro("upscale_factor",
                                                &PixelShuffleAttrs::upscale_factor,
                                                "Scale factor for spatial upsampling.");
  }
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.PixelShuffleAttrs", PixelShuffleAttrs, AttrsNode);
};

}  // namespace relax
}  // namespace tvm

#endif  // TVM_RELAX_ATTRS_NN_H_
