nn.h

Go to the documentation of this file.

/*
 * 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.
 */
 
#ifndef TVM_RELAX_ATTRS_NN_H_
#define TVM_RELAX_ATTRS_NN_H_
 
#include <tvm/relax/expr.h>
 
namespace tvm {
namespace relax {
 
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
 
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
 
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
 
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
 
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
 
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
 
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
 
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
 
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
 
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
 
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
 
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
 
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);
};
 
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);
};
 
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);
};
 
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);
};
 
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
 
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
 
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
 
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
 
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
 
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
 
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
 
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
 
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);
};
 
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_