[tvm] 19/46: [Unity] Relax op: neural networks (#13993)

[tqchen]

This is an automated email from the ASF dual-hosted git repository.

tqchen pushed a commit to branch unity-staging
in repository https://gitbox.apache.org/repos/asf/tvm.git

commit 7dd27f075832f9ea9d89d1116f637de8480bd5c9
Author: Ruihang Lai <ruiha...@cs.cmu.edu>
AuthorDate: Tue Feb 14 15:03:22 2023 -0500

    [Unity] Relax op: neural networks (#13993)
    
    This PR is about the high-level tensor computation operators in Relax.
    
    This PR includes the neural network operators.
---
 include/tvm/relax/attrs/nn.h                       | 190 +++++
 python/tvm/relax/op/__init__.py                    |   1 +
 python/tvm/relax/op/{ => nn}/__init__.py           |  31 +-
 .../tvm/relax/op/{__init__.py => nn/_ffi_api.py}   |  30 +-
 python/tvm/relax/op/nn/nn.py                       | 524 ++++++++++++
 python/tvm/relax/op/op_attrs.py                    |  35 +
 python/tvm/script/ir_builder/relax/ir.py           |   2 +
 src/relax/op/nn/convolution.cc                     | 146 ++++
 src/relax/op/nn/convolution.h                      |  63 ++
 src/relax/op/nn/nn.cc                              | 245 ++++++
 src/relax/op/nn/nn.h                               |  81 ++
 src/relax/op/nn/pooling.cc                         | 184 ++++
 src/relax/op/nn/pooling.h                          |  46 +
 tests/python/relax/test_op_nn.py                   | 929 +++++++++++++++++++++
 tests/python/relax/test_op_nn_convolution.py       | 429 ++++++++++
 tests/python/relax/test_op_nn_pooling.py           | 429 ++++++++++
 tests/python/relax/test_tvmscript_parser_op_nn.py  | 193 +++++
 17 files changed, 3503 insertions(+), 55 deletions(-)

diff --git a/include/tvm/relax/attrs/nn.h b/include/tvm/relax/attrs/nn.h
new file mode 100644
index 0000000000..694a510706
--- /dev/null
+++ b/include/tvm/relax/attrs/nn.h
@@ -0,0 +1,190 @@
+/*
+ * 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 Conv2d operator */
+struct Conv2DAttrs : public tvm::AttrsNode<Conv2DAttrs> {
+  Array<IntImm> strides;
+  Array<IntImm> padding;
+  Array<IntImm> dilation;
+  int groups;
+  String data_layout;
+  String kernel_layout;
+  String out_layout;
+  DataType out_dtype;
+
+  TVM_DECLARE_ATTRS(Conv2DAttrs, "relax.attrs.Conv2DAttrs") {
+    TVM_ATTR_FIELD(strides).describe("Specifies the strides of the 
convolution.");
+    TVM_ATTR_FIELD(padding).describe(
+        "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)");
+    TVM_ATTR_FIELD(dilation).describe(
+        "Specifies the dilation rate to use for dilated convolution.");
+    TVM_ATTR_FIELD(groups).describe(
+        "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.");
+    TVM_ATTR_FIELD(data_layout)
+        .describe(
+            "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.");
+    TVM_ATTR_FIELD(kernel_layout)
+        .describe(
+            "Dimension ordering of weight. Can be 'OIHW', 'OIHW16o16i', etc."
+            "'O', 'I', 'H', 'W' stands for num_filter, input_channel, height, 
and width"
+            "dimensions respectively.");
+    TVM_ATTR_FIELD(out_layout)
+        .describe(
+            "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.");
+    TVM_ATTR_FIELD(out_dtype).describe(
+        "Output data type, set to explicit type under mixed precision 
setting");
+  }
+};  // struct Conv2dAttrs
+
+/*! \brief Attributes used in max_pool2d operator */
+struct MaxPool2DAttrs : public tvm::AttrsNode<MaxPool2DAttrs> {
+  Array<IntImm> pool_size;
+  Array<IntImm> strides;
+  Array<IntImm> padding;
+  Array<IntImm> dilation;
+  bool ceil_mode;
+  String layout;
+  String out_layout;
+
+  TVM_DECLARE_ATTRS(MaxPool2DAttrs, "relax.attrs.MaxPool2DAttrs") {
+    TVM_ATTR_FIELD(pool_size).describe("Size of the pooling windows.");
+    TVM_ATTR_FIELD(strides).describe("Specifies the strides of the 
convolution.");
+    TVM_ATTR_FIELD(dilation).describe("Specifies the dilation of the 
convolution.");
+    TVM_ATTR_FIELD(padding).describe(
+        "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)");
+    TVM_ATTR_FIELD(ceil_mode).describe(
+        "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.");
+    TVM_ATTR_FIELD(layout).describe(
+        "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.");
+    TVM_ATTR_FIELD(out_layout)
+        .describe(
+            "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.");
+  }
+};  // struct MaxPool2dAttrs
+
+/*! \brief Attributes for 2d adaptive pool operator */
+struct AdaptivePool2DAttrs : public tvm::AttrsNode<AdaptivePool2DAttrs> {
+  Optional<Array<IntImm>> output_size;
+  String layout;
+  String out_layout;
+
+  TVM_DECLARE_ATTRS(AdaptivePool2DAttrs, "relax.attrs.AdaptivePool2DAttrs") {
+    TVM_ATTR_FIELD(output_size).describe("Output height and width.");
+    TVM_ATTR_FIELD(layout).describe(
+        "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.");
+    TVM_ATTR_FIELD(out_layout)
+        .describe(
+            "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.");
+  }
+};  // struct AdaptivePool2DAttrs
+
+/*! \brief Attributes used in softmax operators */
+struct SoftmaxAttrs : public tvm::AttrsNode<SoftmaxAttrs> {
+  int axis;
+
+  TVM_DECLARE_ATTRS(SoftmaxAttrs, "relax.attrs.SoftmaxAttrs") {
+    TVM_ATTR_FIELD(axis).describe("The axis to sum over when computing 
softmax.");
+  }
+};
+
+/*! \brief Attributes used in batch_norm operator */
+struct BatchNormAttrs : public tvm::AttrsNode<BatchNormAttrs> {
+  int axis;
+  double epsilon;
+  bool center;
+  bool scale;
+
+  TVM_DECLARE_ATTRS(BatchNormAttrs, "relax.attrs.BatchNormAttrs") {
+    TVM_ATTR_FIELD(axis).describe("The axis along which the normalization is 
applied.");
+    TVM_ATTR_FIELD(epsilon).describe("Small float added to variance to avoid 
dividing by zero");
+    TVM_ATTR_FIELD(center).describe(
+        "Indicating if the beta offset will be added to the normalized 
tensor.");
+    TVM_ATTR_FIELD(scale).describe("Indicating if the gamma scale will be 
multiplied.");
+  }
+};  // struct BatchNormAttrs
+
+/*! \brief Attributes used in layer_norm operator */
+struct LayerNormAttrs : public tvm::AttrsNode<LayerNormAttrs> {
+  Array<Integer> axes;
+  double epsilon;
+  bool center;
+  bool scale;
+
+  TVM_DECLARE_ATTRS(LayerNormAttrs, "relax.attrs.LayerNormAttrs") {
+    TVM_ATTR_FIELD(axes).describe("The axes that along which the normalization 
is applied.");
+    TVM_ATTR_FIELD(epsilon).describe("Small float added to variance to avoid 
dividing by zero");
+    TVM_ATTR_FIELD(center).describe(
+        "Indicating if the beta offset will be added to the normalized 
tensor.");
+    TVM_ATTR_FIELD(scale).describe("Indicating if the gamma scale will be 
multiplied.");
+  }
+};  // struct LayerNormAttrs
+
+/*! \brief Attributes used in dropout operator */
+struct DropoutAttrs : public tvm::AttrsNode<DropoutAttrs> {
+  double rate;
+
+  TVM_DECLARE_ATTRS(DropoutAttrs, "relax.attrs.DropoutAttrs") {
+    TVM_ATTR_FIELD(rate).describe(
+        "Fraction of the input that gets dropped out during training time");
+  }
+};  // struct DropoutAttrs
+
+}  // namespace relax
+}  // namespace tvm
+
+#endif  // TVM_RELAX_ATTRS_NN_H_
diff --git a/python/tvm/relax/op/__init__.py b/python/tvm/relax/op/__init__.py
index 68152c2056..6c6fffc7c6 100644
--- a/python/tvm/relax/op/__init__.py
+++ b/python/tvm/relax/op/__init__.py
@@ -31,6 +31,7 @@ from .unary import *
 from . import builtin
 from . import image
 from . import memory
+from . import nn
 
 
 def _register_op_make():
diff --git a/python/tvm/relax/op/__init__.py 
b/python/tvm/relax/op/nn/__init__.py
similarity index 57%
copy from python/tvm/relax/op/__init__.py
copy to python/tvm/relax/op/nn/__init__.py
index 68152c2056..af2aa106bc 100644
--- a/python/tvm/relax/op/__init__.py
+++ b/python/tvm/relax/op/nn/__init__.py
@@ -14,31 +14,6 @@
 # KIND, either express or implied.  See the License for the
 # specific language governing permissions and limitations
 # under the License.
-# pylint: disable=wildcard-import, redefined-builtin
-"""Relax core operators."""
-
-# Operators
-from .base import *
-from .binary import *
-from .datatype import *
-from .index import *
-from .manipulate import *
-from .op_attrs import *
-from .statistical import *
-from .set import *
-from .ternary import *
-from .unary import *
-from . import builtin
-from . import image
-from . import memory
-
-
-def _register_op_make():
-    # pylint: disable=import-outside-toplevel
-    from . import _ffi_api
-    from .. import expr
-
-    expr._op_ffi_api = _ffi_api  # type: ignore
-
-
-_register_op_make()
+# pylint: disable=wildcard-import
+"""Neural network related operators."""
+from .nn import *
diff --git a/python/tvm/relax/op/__init__.py 
b/python/tvm/relax/op/nn/_ffi_api.py
similarity index 57%
copy from python/tvm/relax/op/__init__.py
copy to python/tvm/relax/op/nn/_ffi_api.py
index 68152c2056..1785345ac1 100644
--- a/python/tvm/relax/op/__init__.py
+++ b/python/tvm/relax/op/nn/_ffi_api.py
@@ -14,31 +14,7 @@
 # KIND, either express or implied.  See the License for the
 # specific language governing permissions and limitations
 # under the License.
-# pylint: disable=wildcard-import, redefined-builtin
-"""Relax core operators."""
+"""Constructor APIs"""
+import tvm._ffi
 
-# Operators
-from .base import *
-from .binary import *
-from .datatype import *
-from .index import *
-from .manipulate import *
-from .op_attrs import *
-from .statistical import *
-from .set import *
-from .ternary import *
-from .unary import *
-from . import builtin
-from . import image
-from . import memory
-
-
-def _register_op_make():
-    # pylint: disable=import-outside-toplevel
-    from . import _ffi_api
-    from .. import expr
-
-    expr._op_ffi_api = _ffi_api  # type: ignore
-
-
-_register_op_make()
+tvm._ffi._init_api("relax.op.nn", __name__)
diff --git a/python/tvm/relax/op/nn/nn.py b/python/tvm/relax/op/nn/nn.py
new file mode 100644
index 0000000000..cdf0e96464
--- /dev/null
+++ b/python/tvm/relax/op/nn/nn.py
@@ -0,0 +1,524 @@
+# 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.
+"""Relax Neural Network (NN) operators"""
+from typing import List, Optional, Tuple, Union
+
+from tvm import DataType
+
+from . import _ffi_api
+from ...expr import Expr
+
+
+def conv2d(
+    data: Expr,
+    weight: Expr,
+    strides: Union[int, Tuple[int, int]] = (1, 1),
+    padding: Union[int, Tuple[int, ...]] = (0, 0),
+    dilation: Union[int, Tuple[int, int]] = (1, 1),
+    groups: int = 1,
+    data_layout: str = "NCHW",
+    kernel_layout: str = "OIHW",
+    out_layout: Optional[str] = None,
+    out_dtype: Optional[Union[str, DataType]] = None,
+) -> Expr:
+    r"""2D convolution.
+
+    This operator takes the weight as the convolution kernel
+    and convolves it with data to produce an output.
+
+
+    In the default case, where the data_layout is `NCHW`
+    and kernel_layout is `OIHW`, conv2d takes in
+    a data Tensor with shape `(batch_size, in_channels, height, width)`,
+    and a weight Tensor with shape `(channels, in_channels, kernel_h, 
kernel_w)`,
+    where `kernel_h` and `kernel_w` is the lengths of the `H` and `W` kernel 
dimensions,
+    to produce an output Tensor with the following rule:
+
+    .. math::
+
+        \mbox{out}[b, c, y, x] = \sum_{dy, dx, k}
+           \mbox{data}[b, k, \mbox{strides}[0] * y  + dy, \mbox{strides}[1] * 
x + dx] *
+           \mbox{weight}[c, k, dy, dx]
+
+    Padding and dilation are applied to data and weight respectively before 
the computation.
+    This operator accepts data layout specification.
+    Semantically, the operator will convert the layout to the canonical layout
+    (`NCHW` for data and `OIHW` for weight), perform the computation,
+    then convert to the out_layout.
+
+    Parameters
+    ----------
+    data : relax.Expr
+        The input data to the operator.
+
+    weight : relax.Expr
+        The weight expressions.
+
+    strides : Union[int, Tuple[int, int]]
+        The strides of convolution. It is required to have length either 1 or 
2.
+
+    padding : Union[int, Tuple[int, ...]]
+        The padding of convolution on both sides of inputs before convolution.
+        It is required to have length either 1, 2 or 4.
+
+    dilation : Union[int, Tuple[int, int]]
+        Specifies the dilation rate to be used for dilated convolution.
+        It is required to have length either 1 or 2.
+
+    groups : int
+        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.
+
+    data_layout : str
+        Layout of the input.
+
+    kernel_layout : str
+        Layout of the weight.
+
+    out_layout : Optional[str]
+        Layout of the output. If not specified, it is the same as data_layout
+
+    out_dtype : Optional[Union[str, DataType]]
+        Specifies the output data type for mixed precision conv2d.
+
+    Returns
+    -------
+    result : relax.Expr
+        The computed result.
+    """
+    if isinstance(strides, int):
+        strides = (strides, strides)
+    if isinstance(dilation, int):
+        dilation = (dilation, dilation)
+    if isinstance(padding, int):
+        padding = (padding, padding, padding, padding)
+
+    return _ffi_api.conv2d(  # type: ignore
+        data,
+        weight,
+        strides,
+        padding,
+        dilation,
+        groups,
+        data_layout,
+        kernel_layout,
+        out_layout,
+        out_dtype,
+    )
+
+
+def max_pool2d(
+    data: Expr,
+    pool_size: Union[int, Tuple[int, int]] = (1, 1),
+    strides: Union[int, Tuple[int, int]] = (1, 1),
+    padding: Union[int, Tuple[int, ...]] = (0, 0),
+    dilation: Union[int, Tuple[int, int]] = (1, 1),
+    ceil_mode: bool = False,
+    layout: str = "NCHW",
+    out_layout: Optional[str] = None,
+) -> Expr:
+    r"""2D maximum pooling operator.
+
+    This operator takes data as input and does 2D max value calculation
+    with in pool_size sized window by striding defined by stride
+
+
+    In the default case, where the data_layout is `NCHW`
+    a data Tensor with shape `(batch_size, in_channels, height, width)`,
+    to produce an output Tensor with the following rule:
+
+    with data of shape (b, c, h, w) and pool_size (kh, kw)
+
+    .. math::
+
+        \mbox{out}(b, c, y, x)  = \max_{m=0, \ldots, kh-1} \max_{n=0, \ldots, 
kw-1}
+             \mbox{data}(b, c, \mbox{stride}[0] * y + m, \mbox{stride}[1] * x 
+ n)
+
+    Padding is applied to data before the computation.
+    ceil_mode is used to take ceil or floor while computing out shape.
+    This operator accepts data layout specification.
+
+    Parameters
+    ----------
+    data : relax.Expr
+        The input data to the operator.
+
+    pool_size : Union[int, Tuple[int, int]]
+        The size of window for pooling. It is required to have length either 1 
or 2.
+
+    strides : Union[int, Tuple[int, int]]
+        The strides of pooling. It is required to have length either 1 or 2.
+
+    padding : Union[int, Tuple[int, ...]]
+        The padding for pooling. It is required to have length either 1, 2 or 
4.
+
+    dilation : Union[int, Tuple[int, int]]
+        The dilation of pooling. It is required to have length either 1 or 2.
+
+    ceil_mode : bool
+        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.
+
+    layout : str
+        Layout of the input.
+
+    out_layout : Optional[str]
+        Layout of the output. If not specified, it is the same as data_layout
+
+    Returns
+    -------
+    result : Expr
+        The computed result.
+    """
+    if isinstance(pool_size, int):
+        pool_size = (pool_size, pool_size)
+    if isinstance(strides, int):
+        strides = (strides, strides)
+    if isinstance(dilation, int):
+        dilation = (dilation, dilation)
+    if isinstance(padding, int):
+        padding = (padding, padding, padding, padding)
+
+    return _ffi_api.max_pool2d(  # type: ignore
+        data, pool_size, strides, padding, dilation, ceil_mode, layout, 
out_layout
+    )
+
+
+def adaptive_avg_pool2d(
+    data: Expr,
+    output_size: Optional[Union[int, Tuple[int, int]]] = None,
+    layout: str = "NCHW",
+    out_layout: Optional[str] = None,
+) -> Expr:
+    r"""2D adaptive average pooling operator. This operator is experimental.
+
+    This operator takes data as input and does 2D average value calculation
+    across each window represented by WxH.
+
+
+    In the default case, where the data_layout is `NCHW`
+    a data Tensor with shape `(batch_size, in_channels, height, width)`,
+    to produce an output Tensor with shape
+    (batch_size, in_channels, output_height, output_width).
+
+    The pooling kernel and stride sizes are automatically chosen for
+    desired output sizes.
+
+    For output_size:
+        If this argument is not provided, input height and width will be used
+        as output height and width.
+
+        If a single integer is provided for output_size, the output size is
+        (N x C x output_size x output_size) for any input (NCHW).
+
+        If a tuple of integers (height, width) are provided for output_size,
+        the output size is (N x C x height x width) for any input (NCHW).
+
+    Parameters
+    ----------
+    data : relax.Expr
+        The input data to the operator.
+
+    output_size : Optional[Union[int, Tuple[int, int]]]
+        Output height and width.
+        If not specified, it will be the same as the input height and width.
+        If specified, it is required to have length either 1 or 2.
+
+    layout : str
+        Layout of the input.
+
+    out_layout : Optional[str]
+        Layout of the output. If not specified, it is the same as data_layout
+
+    Returns
+    -------
+    result : relax.Expr
+        The computed result.
+    """
+    if isinstance(output_size, int):
+        output_size = (output_size, output_size)
+    return _ffi_api.adaptive_avg_pool2d(data, output_size, layout, out_layout) 
 # type: ignore
+
+
+def relu(data: Expr) -> Expr:
+    """Rectified linear unit.
+
+    .. math::
+        text{ReLU}(x) = max(x, 0)
+
+    Parameters
+    ----------
+    data : relax.Expr
+        The input data
+
+    Returns
+    -------
+    result : relax.Expr
+        The computed result.
+    """
+    return _ffi_api.relu(data)  # type: ignore
+
+
+def gelu(data: Expr) -> Expr:
+    """Gaussian Error Linear Units function
+
+    .. math::
+        text{GeLU}(x) = 0.5 * x * (1 + erf(x * 0.5**0.5))
+
+    where :math:`erf` is the Gauss Error function.
+
+    Parameters
+    ----------
+    data : relax.Expr
+        The input data
+
+    Returns
+    -------
+    result : relax.Expr
+        The computed result.
+
+    Note
+    ----
+    The input tensor is required to have float dtype
+    """
+    return _ffi_api.gelu(data)  # type: ignore
+
+
+def silu(data: Expr) -> Expr:
+    """Sigmoid Linear Unit function
+
+    .. math::
+        text{SiLU}(x) = x * sigmoid(x)
+
+    Parameters
+    ----------
+    data : relax.Expr
+        The input data
+
+    Returns
+    -------
+    result : relax.Expr
+        The computed result.
+
+    Note
+    ----
+    The input tensor is required to have float dtype
+    """
+    return _ffi_api.silu(data)  # type: ignore
+
+
+def softmax(data: Expr, axis: int = -1) -> Expr:
+    r"""Computes softmax.
+
+    .. math:: text{softmax}(x)_i = frac{exp(x_i)}{\sum_j exp(x_j)}
+
+    Parameters
+    ----------
+    data: relax.Expr
+        The input data to the operator.
+
+    axis: int
+        The axis to sum over when computing softmax.
+        If not specified, it is by default the last axis of the input tensor.
+        Supports negative indexing.
+
+    Returns
+    -------
+    result : relax.Expr
+        The computed result.
+
+    Note
+    ----
+    The input tensor is required to have float dtype
+    """
+    return _ffi_api.softmax(data, axis)  # type: ignore
+
+
+def batch_norm(
+    data: Expr,
+    gamma: Expr,
+    beta: Expr,
+    moving_mean: Expr,
+    moving_var: Expr,
+    axis: int,
+    epsilon: float = 1e-5,
+    center: bool = True,
+    scale: bool = True,
+) -> Expr:
+    r"""
+    Batch normalization layer (Ioffe and Szegedy, 2014).
+    Normalizes the input at each batch, i.e. applies a transformation
+    that maintains the mean activation close to 0 and the activation
+    standard deviation close to 1.
+
+    .. math::
+
+        data\_mean[i] = mean(data[:,i,:,...]) \\
+        data\_var[i] = var(data[:,i,:,...])
+
+    Then compute the normalized output, which has the same shape as input, as 
following:
+
+    .. math::
+
+        out[:,i,:,...] = \frac{data[:,i,:,...] - 
data\_mean[i]}{\sqrt{data\_var[i]+\epsilon}}
+            * gamma[i] + beta[i]
+
+    Both *mean* and *var* returns a scalar by treating the input as a vector.
+
+    Assume the input has size *k* on axis 1, then both ``gamma`` and ``beta``
+    have shape *(k,)*.
+
+    Besides the inputs and the outputs, this operator accepts two auxiliary
+    states, ``moving_mean`` and ``moving_var``, which are *k*-length
+    vectors. They are global statistics for the whole dataset, which are 
updated by
+
+    .. code:: python
+
+        moving_mean = moving_mean * momentum + data_mean * (1 - momentum)
+        moving_var = moving_var * momentum + data_var * (1 - momentum)
+
+    The parameter ``axis`` specifies which axis of the input shape denotes
+    the 'channel' (separately normalized groups).  The default is 1.
+    Specifying -1 sets the channel axis to be the last item in the input shape.
+
+    .. note::
+
+        This operator can be optimized away for inference.
+
+    Parameters
+    ----------
+    data : relax.Expr
+        The input data to the operator.
+
+    gamma : relax.Expr
+        The gamma scale factor.
+
+    beta : relax.Expr
+        The beta offset factor.
+
+    moving_mean : relax.Expr
+        Running mean of input.
+
+    moving_var : relax.Expr
+        Running variance of input.
+
+    axis : int
+        The axis along which the normalization is applied.
+
+    epsilon : float
+        Small float added to variance to avoid dividing by zero.
+
+    center : bool
+        Indicating if the beta offset will be added to the normalized tensor.
+
+    scale : bool
+        Indicating if the gamma scale will be multiplied.
+
+    Returns
+    -------
+    result : relax.Expr
+        The computed result.
+    """
+    return _ffi_api.batch_norm(  # type: ignore
+        data, gamma, beta, moving_mean, moving_var, axis, epsilon, center, 
scale
+    )
+
+
+def layer_norm(
+    data: Expr,
+    gamma: Expr,
+    beta: Expr,
+    axes: Union[int, List[int]],
+    epsilon: float = 1e-5,
+    center: bool = True,
+    scale: bool = True,
+) -> Expr:
+    r"""
+    Layer normalization (Lei Ba and et al., 2016).
+    Applies layer normalization to the n-dimensional input array.
+    This operator takes an n-dimensional input array and normalizes
+    the input using the given axis:
+
+    .. math::
+
+        out = \frac{data - mean(data, axis)}{\sqrt{var(data, axis)+\epsilon}}
+            * gamma + beta
+
+    Unlike batch normalization, the mean and var are computed along the 
channel dimension.
+
+    Assume the input has size k on axis 1, then both gamma and beta have shape 
(k,).
+
+    .. note::
+
+        This operator can be optimized away for inference.
+
+    Parameters
+    ----------
+    data : relax.Expr
+        Input to which layer_norm will be applied.
+
+    gamma : relax.Expr
+        The gamma scale factor.
+
+    beta : relax.Expr
+        The beta offset factor.
+
+    axes : Union[int, List[int]]
+        The axes that along which the normalization is applied.
+
+    epsilon : float
+        Small float added to variance to avoid dividing by zero.
+
+    center : bool
+        Indicating if the beta offset will be added to the normalized tensor.
+
+    scale : bool
+        Indicating if the gamma scale will be multiplied.
+
+    Returns
+    -------
+    result : relax.Expr
+        The computed result.
+    """
+    if isinstance(axes, int):
+        axes = [axes]
+    return _ffi_api.layer_norm(data, gamma, beta, axes, epsilon, center, 
scale)  # type: ignore
+
+
+def dropout(data: Expr, rate: float = 0.5) -> Expr:
+    """Applies the dropout operation to the input tensor.
+
+    During training, each element of the input is set to zero with
+    probability ``p``. The whole array is scaled by ``1/(1-p)``
+    to keep the expected sum of the input unchanged.
+
+    Parameters
+    ----------
+    data : relax.Expr
+        The input data to the operator.
+
+    rate : float
+        The probability for an element to be reset to 0.
+
+    Returns
+    -------
+    result : relax.Expr
+        The result of dropout, which is a tuple of two tensors.
+        The first one is the original tensor and the second one is a
+        mask tensor (1.0 where element not dropped, 0.0 where dropped)
+    """
+    return _ffi_api.dropout(data, rate)  # type: ignore
diff --git a/python/tvm/relax/op/op_attrs.py b/python/tvm/relax/op/op_attrs.py
index 1fb8853040..68f84b3514 100644
--- a/python/tvm/relax/op/op_attrs.py
+++ b/python/tvm/relax/op/op_attrs.py
@@ -34,6 +34,41 @@ class StridedSliceAttrs(Attrs):
     """Attributes used in strided_slice operator"""
 
 
+@tvm._ffi.register_object("relax.attrs.Conv2DAttrs")
+class Conv2DAttrs(Attrs):
+    """Attributes for nn.conv2d"""
+
+
+@tvm._ffi.register_object("relax.attrs.MaxPool2DAttrs")
+class MaxPool2DAttrs(Attrs):
+    """Attributes for nn.max_pool2d"""
+
+
+@tvm._ffi.register_object("relax.attrs.AdaptivePool2DAttrs")
+class AdaptivePool2DAttrs(Attrs):
+    """Attributes for 2d adaptive pool operator"""
+
+
+@tvm._ffi.register_object("relax.attrs.SoftmaxAttrs")
+class SoftmaxAttrs(Attrs):
+    """Attributes for nn.softmax"""
+
+
+@tvm._ffi.register_object("relax.attrs.BatchNormAttrs")
+class BatchNormAttrs(Attrs):
+    """Attributes used in batch_norm operator"""
+
+
+@tvm._ffi.register_object("relax.attrs.LayerNormAttrs")
+class LayerNormAttrs(Attrs):
+    """Attributes used in layer_norm operator"""
+
+
+@tvm._ffi.register_object("relax.attrs.DropoutAttrs")
+class DropoutAttrs(Attrs):
+    """Attributes for dropout operator"""
+
+
 @tvm._ffi.register_object("relax.attrs.StatisticalAttrs")
 class StatisticalAttrs(Attrs):
     """Attributes used in statistical operator"""
diff --git a/python/tvm/script/ir_builder/relax/ir.py 
b/python/tvm/script/ir_builder/relax/ir.py
index 47779a6024..1f0e31428c 100644
--- a/python/tvm/script/ir_builder/relax/ir.py
+++ b/python/tvm/script/ir_builder/relax/ir.py
@@ -93,6 +93,7 @@ from tvm.relax.op import (
     tan,
     tanh,
     unique,
+    nn,
 )
 from tvm.relax.struct_info import StructInfo
 from tvm.relax.utils import args_converter
@@ -530,4 +531,5 @@ __all__ = [
     "tuple",
     "variance",
     "unique",
+    "nn",    
 ]
diff --git a/src/relax/op/nn/convolution.cc b/src/relax/op/nn/convolution.cc
new file mode 100644
index 0000000000..a3ddd3e350
--- /dev/null
+++ b/src/relax/op/nn/convolution.cc
@@ -0,0 +1,146 @@
+/*
+ * 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 src/relax/op/nn/convolution.cc
+ * \brief Convolution operators
+ */
+
+#include "convolution.h"
+
+#include <vector>
+
+namespace tvm {
+namespace relax {
+
+/* relax.nn.conv2d */
+TVM_REGISTER_NODE_TYPE(Conv2DAttrs);
+
+Expr conv2d(Expr data, Expr weight, Array<IntImm> strides, Array<IntImm> 
padding,
+            Array<IntImm> dilation, int groups, String data_layout, String 
kernel_layout,
+            Optional<String> out_layout, DataType out_dtype) {
+  padding = GetCompletePadding2D(std::move(padding));
+  if (strides.size() == 1) {
+    strides.push_back(strides[0]);
+  }
+  if (dilation.size() == 1) {
+    dilation.push_back(dilation[0]);
+  }
+
+  CHECK_GT(groups, 0) << "The number of groups in convolution is expected to 
be positive. However, "
+                         "the given number of groups is "
+                      << groups;
+  CHECK_EQ(strides.size(), 2)
+      << "The input strides length is expected to be 2. However, the given 
strides is " << strides;
+  CHECK_EQ(dilation.size(), 2)
+      << "The input dilation length is expected to be 2. However, the given 
dilation is "
+      << dilation;
+  return MakeConv<Conv2DAttrs>(std::move(data), std::move(weight), 
std::move(strides),
+                               std::move(padding), std::move(dilation), 
groups, data_layout,
+                               std::move(kernel_layout), 
out_layout.value_or(data_layout),
+                               out_dtype, /*op_name=*/"relax.nn.conv2d");
+}
+
+TVM_REGISTER_GLOBAL("relax.op.nn.conv2d").set_body_typed(conv2d);
+
+StructInfo InferStructInfoConv2d(const Call& call, const BlockBuilder& ctx) {
+  Array<TensorStructInfo> input_sinfo = GetInputTensorStructInfo(call, ctx);
+  TensorStructInfo data_sinfo = input_sinfo[0];
+  TensorStructInfo weight_sinfo = input_sinfo[1];
+
+  const auto* attrs = call->attrs.as<Conv2DAttrs>();
+  auto [data_layout, data2NCHW] = CheckTensorLayout(call, ctx, 
attrs->data_layout,  //
+                                                    /*tgt_layout=*/"NCHW",     
     //
+                                                    /*tensor_name=*/"data");
+  auto [weight_layout, weight2OIHW] = CheckTensorLayout(call, ctx, 
attrs->kernel_layout,  //
+                                                        /*tgt_layout=*/"OIHW", 
           //
+                                                        
/*tensor_name=*/"kernel");
+  auto [out_layout, out2NCHW] = CheckTensorLayout(call, ctx, 
attrs->out_layout,  //
+                                                  /*tgt_layout=*/"NCHW",       
  //
+                                                  /*tensor_name=*/"output");
+
+  Optional<ShapeExpr> data_shape =
+      CheckNdimPerLayoutAndGetShape(call, ctx, data_sinfo, data_layout);
+  Optional<ShapeExpr> weight_shape =
+      CheckNdimPerLayoutAndGetShape(call, ctx, weight_sinfo, weight_layout);
+
+  DataType out_dtype = attrs->out_dtype.is_void()
+                           ? InferBinaryArithOpOutDtype(call, ctx, data_sinfo, 
weight_sinfo)
+                           : attrs->out_dtype;
+  if (!data_shape.defined() || !weight_shape.defined()) {
+    return TensorStructInfo(out_dtype, out_layout.ndim());
+  }
+
+  Array<PrimExpr> data_NCHW_shape = 
data2NCHW.ForwardShape(data_shape.value()->values);
+  Array<PrimExpr> weight_OIHW_shape = 
weight2OIHW.ForwardShape(weight_shape.value()->values);
+
+  arith::Analyzer* analyzer = ctx->GetAnalyzer();
+  PrimExpr input_channel_data = data_NCHW_shape[1];
+  PrimExpr input_channel_kernel = weight_OIHW_shape[1];
+  if (analyzer->CanProve(input_channel_data != input_channel_kernel * 
attrs->groups)) {
+    ctx->ReportFatal(
+        Diagnostic::Error(call)
+        << "The channel size of the data should equal to the product of input 
channel size of the "
+           "weight and the number of groups. However, the data channel size is 
"
+        << input_channel_data << " while the weight input channel size and 
number of groups are "
+        << input_channel_kernel << " and " << attrs->groups);
+  } else if (!analyzer->CanProveEqual(input_channel_data, input_channel_kernel 
* attrs->groups)) {
+    // Todo(relax-team): Trust the input shape at this moment, and revisit
+    // this condition with runtime shape check
+  }
+  if (analyzer->CanProve(floormod(weight_OIHW_shape[0], attrs->groups) != 0)) {
+    ctx->ReportFatal(Diagnostic::Error(call)
+                     << "Conv2d expects the number of output channels to be 
divisible by the "
+                        "number of groups. However, the number of output 
channels is "
+                     << weight_OIHW_shape[0] << " while the number of groups 
is " << attrs->groups);
+  } else if (!analyzer->CanProveEqual(floormod(weight_OIHW_shape[0], 
attrs->groups), 0)) {
+    // Todo(relax-team): Trust the input shape at this moment, and revisit
+    // this condition with runtime shape check
+  }
+
+  PrimExpr input_h = data_NCHW_shape[2];
+  PrimExpr input_w = data_NCHW_shape[3];
+  PrimExpr kernel_h = weight_OIHW_shape[2];
+  PrimExpr kernel_w = weight_OIHW_shape[3];
+  PrimExpr padding_h = attrs->padding[0] + attrs->padding[2];
+  PrimExpr padding_w = attrs->padding[1] + attrs->padding[3];
+
+  std::vector<PrimExpr> out_NCHW_shape;
+  out_NCHW_shape.resize(4);
+  out_NCHW_shape[0] = data_NCHW_shape[0];
+  out_NCHW_shape[1] = weight_OIHW_shape[0];
+
+  PrimExpr numerator_h = input_h + padding_h - attrs->dilation[0] * (kernel_h 
- 1) - 1;
+  PrimExpr numerator_w = input_w + padding_w - attrs->dilation[1] * (kernel_w 
- 1) - 1;
+  out_NCHW_shape[2] = analyzer->Simplify(floordiv(numerator_h, 
attrs->strides[0]) + 1);
+  out_NCHW_shape[3] = analyzer->Simplify(floordiv(numerator_w, 
attrs->strides[1]) + 1);
+
+  Array<PrimExpr> out_shape = out2NCHW.BackwardShape(out_NCHW_shape);
+  return TensorStructInfo(ShapeExpr(out_shape), out_dtype);
+}
+
+TVM_REGISTER_OP("relax.nn.conv2d")
+    .set_num_inputs(2)
+    .add_argument("data", "Tensor", "The input tensor.")
+    .add_argument("weight", "Tensor", "The weight tensor.")
+    .set_attrs_type<Conv2DAttrs>()
+    .set_attr<FInferStructInfo>("FInferStructInfo", InferStructInfoConv2d);
+
+}  // namespace relax
+}  // namespace tvm
diff --git a/src/relax/op/nn/convolution.h b/src/relax/op/nn/convolution.h
new file mode 100644
index 0000000000..a65617b48d
--- /dev/null
+++ b/src/relax/op/nn/convolution.h
@@ -0,0 +1,63 @@
+/*
+ * 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 convolution.h
+ * \brief The functions to make Relax neural network convolution operator 
calls.
+ */
+
+#ifndef TVM_RELAX_OP_NN_CONVOLUTION_H_
+#define TVM_RELAX_OP_NN_CONVOLUTION_H_
+
+#include <tvm/relax/attrs/nn.h>
+
+#include <string>
+#include <utility>
+
+#include "../op_common.h"
+
+namespace tvm {
+namespace relax {
+
+template <typename T>
+inline Expr MakeConv(Expr data, Expr weight, Array<IntImm> strides, 
Array<IntImm> padding,
+                     Array<IntImm> dilation, int groups, String data_layout, 
String kernel_layout,
+                     String out_layout, DataType out_dtype, std::string 
op_name) {
+  auto attrs = make_object<T>();
+  attrs->strides = ConvertIntImmToInt64(strides);
+  attrs->padding = ConvertIntImmToInt64(padding);
+  attrs->dilation = ConvertIntImmToInt64(dilation);
+  attrs->groups = groups;
+  attrs->data_layout = std::move(data_layout);
+  attrs->kernel_layout = std::move(kernel_layout);
+  attrs->out_layout = std::move(out_layout);
+  attrs->out_dtype = std::move(out_dtype);
+  const Op& op = Op::Get(op_name);
+  return Call(op, {data, weight}, Attrs(attrs), {});
+}
+
+/*! \brief 2D convolution */
+Expr conv2d(Expr data, Expr weight, Array<IntImm> strides, Array<IntImm> 
padding,
+            Array<IntImm> dilation, int groups, String data_layout, String 
kernel_layout,
+            Optional<String> out_layout, DataType out_dtype);
+
+}  // namespace relax
+}  // namespace tvm
+
+#endif  // TVM_RELAX_OP_NN_CONVOLUTION_H_
diff --git a/src/relax/op/nn/nn.cc b/src/relax/op/nn/nn.cc
new file mode 100644
index 0000000000..66ae10fe6c
--- /dev/null
+++ b/src/relax/op/nn/nn.cc
@@ -0,0 +1,245 @@
+/*
+ * 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.
+ */
+
+#include "nn.h"
+
+#include <utility>
+#include <vector>
+
+namespace tvm {
+namespace relax {
+
+/* relax.nn.relu */
+RELAX_REGISTER_UNARY_NN_OP_AND_IMPL(relu, "nn.relu", 
/*require_float_dtype=*/false);
+
+/* relax.nn.gelu */
+RELAX_REGISTER_UNARY_NN_OP_AND_IMPL(gelu, "nn.gelu", 
/*require_float_dtype=*/true);
+
+/* relax.nn.silu */
+RELAX_REGISTER_UNARY_NN_OP_AND_IMPL(silu, "nn.silu", 
/*require_float_dtype=*/true);
+
+/* relax.nn.softmax */
+TVM_REGISTER_NODE_TYPE(SoftmaxAttrs);
+
+Expr softmax(Expr data, int axis) {
+  auto attrs = make_object<SoftmaxAttrs>();
+  attrs->axis = axis;
+  static const Op& op = Op::Get("relax.nn.softmax");
+  return Call(op, {data}, Attrs(attrs), {});
+}
+
+TVM_REGISTER_GLOBAL("relax.op.nn.softmax").set_body_typed(softmax);
+
+StructInfo InferStructInfoSoftmax(const Call& call, const BlockBuilder& ctx) {
+  TensorStructInfo data_sinfo = GetUnaryInputTensorStructInfo(call, ctx);
+  if (data_sinfo->IsUnknownNdim()) {
+    return data_sinfo;
+  }
+  if (!data_sinfo->IsUnknownDtype() && !data_sinfo->dtype.is_float()) {
+    ctx->ReportFatal(Diagnostic::Error(call) << "Softmax requires the input 
tensor to have float "
+                                                "dtype. However, the given 
input dtype is "
+                                             << data_sinfo->dtype);
+  }
+  const auto* attrs = call->attrs.as<SoftmaxAttrs>();
+  NormalizeAxis(call, ctx, data_sinfo->ndim, attrs->axis);
+
+  return data_sinfo;
+}
+
+TVM_REGISTER_OP("relax.nn.softmax")
+    .set_num_inputs(1)
+    .add_argument("data", "Tensor", "The input tensor.")
+    .set_attrs_type<SoftmaxAttrs>()
+    .set_attr<FInferStructInfo>("FInferStructInfo", InferStructInfoSoftmax);
+
+bool NormCheckDtypeAndShape(const Call& call, const BlockBuilder& ctx,
+                            const Array<TensorStructInfo>& input_sinfo, 
Array<Integer> axes) {
+  Op op = Downcast<Op>(call->op);
+  int n_input = op->arguments.size();
+
+  TensorStructInfo data_sinfo = input_sinfo[0];
+
+  std::vector<int> axes_non_neg;
+  if (!data_sinfo->IsUnknownNdim()) {
+    axes_non_neg = NormalizeAxes(call, ctx, data_sinfo->ndim, axes);
+  }
+  int n_axis = axes.size();
+  if (!data_sinfo->IsUnknownDtype() && !data_sinfo->dtype.is_float()) {
+    ctx->ReportFatal(
+        Diagnostic::Error(call)
+        << op << " requires the input data to have float dtype. However, the 
given data dtype is "
+        << data_sinfo->dtype);
+  }
+  for (int i = 1; i < n_input; ++i) {
+    if (input_sinfo[i]->dtype != data_sinfo->dtype) {
+      ctx->ReportFatal(Diagnostic::Error(call)
+                       << op
+                       << " requires all the input tensors to have the same 
dtype. However, the "
+                       << op->arguments[i]->name << " has dtype " << 
input_sinfo[i]->dtype
+                       << " which is other than the input data's dtype " << 
data_sinfo->dtype);
+    } else if (input_sinfo[i]->ndim != n_axis) {
+      ctx->ReportFatal(Diagnostic::Error(call)
+                       << op << " requires the input " << 
op->arguments[i]->name
+                       << " to have as many dimensions as the length of input 
axes. However, the "
+                          "given one has ndim "
+                       << input_sinfo[i]->ndim << ", which is other than the 
length of axes "
+                       << n_axis);
+    }
+  }
+
+  std::vector<Array<PrimExpr>> axis_lengths;
+  axis_lengths.reserve(n_input);
+  if (const auto* data_shape = data_sinfo->shape.as<ShapeExprNode>()) {
+    std::vector<PrimExpr> lengths;
+    lengths.reserve(n_axis);
+    for (int d = 0; d < n_axis; ++d) {
+      lengths.push_back(data_shape->values[axes_non_neg[d]]);
+    }
+    axis_lengths.push_back(lengths);
+  }
+  for (int i = 1; i < n_input; ++i) {
+    if (const auto* shape = input_sinfo[i]->shape.as<ShapeExprNode>()) {
+      axis_lengths.push_back(shape->values);
+    }
+  }
+
+  arith::Analyzer* analyzer = ctx->GetAnalyzer();
+  for (int i = 1; i < static_cast<int>(axis_lengths.size()); ++i) {
+    for (int d = 0; d < n_axis; ++d) {
+      if (analyzer->CanProve(axis_lengths[0][d] != axis_lengths[i][d])) {
+        ctx->ReportFatal(Diagnostic::Error(call)
+                         << op
+                         << " requires the input gamma, beta, etc., to have 
size same as the "
+                            "lengths of the data on the given axes. However, 
there exists "
+                         << axis_lengths[0] << " and " << axis_lengths[i] << " 
that are unequal.");
+      } else if (!analyzer->CanProveEqual(axis_lengths[0][d], 
axis_lengths[i][d])) {
+        return true;
+      }
+    }
+  }
+  return false;
+}
+
+/* relax.nn.batch_norm */
+TVM_REGISTER_NODE_TYPE(BatchNormAttrs);
+
+Expr batch_norm(Expr data, Expr gamma, Expr beta, Expr moving_mean, Expr 
moving_var,  //
+                int axis, double epsilon, bool center, bool scale) {
+  ObjectPtr<BatchNormAttrs> attrs = make_object<BatchNormAttrs>();
+  attrs->axis = axis;
+  attrs->epsilon = epsilon;
+  attrs->center = center;
+  attrs->scale = scale;
+
+  static const Op& op = Op::Get("relax.nn.batch_norm");
+  return Call(op,
+              {std::move(data), std::move(gamma), std::move(beta), 
std::move(moving_mean),
+               std::move(moving_var)},
+              Attrs{attrs}, {});
+}
+
+TVM_REGISTER_GLOBAL("relax.op.nn.batch_norm").set_body_typed(batch_norm);
+
+StructInfo InferStructInfoBatchNorm(const Call& call, const BlockBuilder& ctx) 
{
+  Array<TensorStructInfo> input_sinfo = GetInputTensorStructInfo(call, ctx);
+
+  const auto* attrs = call->attrs.as<BatchNormAttrs>();
+  bool unknown_shape = NormCheckDtypeAndShape(call, ctx, input_sinfo, 
{attrs->axis});
+
+  DataType dtype = input_sinfo[0]->dtype;
+  if (unknown_shape) {
+    return TupleStructInfo({TensorStructInfo(dtype, input_sinfo[0]->ndim),
+                            TensorStructInfo(dtype, /*ndim=*/1),
+                            TensorStructInfo(dtype, /*ndim=*/1)});
+  } else {
+    return TupleStructInfo({input_sinfo[0], input_sinfo[3], input_sinfo[4]});
+  }
+}
+
+TVM_REGISTER_OP("relax.nn.batch_norm")
+    .set_attrs_type<BatchNormAttrs>()
+    .set_num_inputs(5)
+    .add_argument("data", "Tensor", "Input to which batch_norm will be 
applied.")
+    .add_argument("gamma", "Tensor", "The gamma scale factor.")
+    .add_argument("beta", "Tensor", "The beta offset factor.")
+    .add_argument("moving_mean", "Tensor", "Running mean of input.")
+    .add_argument("moving_var", "Tensor", "Running variance of input.")
+    .set_attr<FInferStructInfo>("FInferStructInfo", InferStructInfoBatchNorm);
+
+/* relax.nn.layer_norm */
+TVM_REGISTER_NODE_TYPE(LayerNormAttrs);
+
+Expr layer_norm(Expr data, Expr gamma, Expr beta, Array<Integer> axes, double 
epsilon, bool center,
+                bool scale) {
+  ObjectPtr<LayerNormAttrs> attrs = make_object<LayerNormAttrs>();
+  attrs->axes = std::move(axes);
+  attrs->epsilon = epsilon;
+  attrs->center = center;
+  attrs->scale = scale;
+
+  static const Op& op = Op::Get("relax.nn.layer_norm");
+  return Call(op, {std::move(data), std::move(gamma), std::move(beta)}, 
Attrs{attrs}, {});
+}
+
+TVM_REGISTER_GLOBAL("relax.op.nn.layer_norm").set_body_typed(layer_norm);
+
+StructInfo InferStructInfoLayerNorm(const Call& call, const BlockBuilder& ctx) 
{
+  Array<TensorStructInfo> input_sinfo = GetInputTensorStructInfo(call, ctx);
+
+  const auto* attrs = call->attrs.as<LayerNormAttrs>();
+  bool unknown_shape = NormCheckDtypeAndShape(call, ctx, input_sinfo, 
attrs->axes);
+
+  return unknown_shape ? TensorStructInfo(input_sinfo[0]->dtype, 
input_sinfo[0]->ndim)
+                       : input_sinfo[0];
+}
+
+TVM_REGISTER_OP("relax.nn.layer_norm")
+    .set_attrs_type<LayerNormAttrs>()
+    .set_num_inputs(3)
+    .add_argument("data", "Tensor", "Input to which batch_norm will be 
applied.")
+    .add_argument("gamma", "Tensor", "The gamma scale factor.")
+    .add_argument("beta", "Tensor", "The beta offset factor.")
+    .set_attr<FInferStructInfo>("FInferStructInfo", InferStructInfoLayerNorm);
+
+/* relax.nn.dropout */
+TVM_REGISTER_NODE_TYPE(DropoutAttrs);
+
+Expr dropout(Expr data, double rate) {
+  ObjectPtr<DropoutAttrs> attrs = make_object<DropoutAttrs>();
+  attrs->rate = rate;
+
+  static const Op& op = Op::Get("relax.nn.dropout");
+  return Call(op, {std::move(data)}, Attrs{attrs}, {});
+}
+
+TVM_REGISTER_GLOBAL("relax.op.nn.dropout").set_body_typed(dropout);
+
+StructInfo InferStructInfoDropout(const Call& call, const BlockBuilder& ctx) {
+  TensorStructInfo data_sinfo = GetUnaryInputTensorStructInfo(call, ctx);
+  return TupleStructInfo({data_sinfo, data_sinfo});
+}
+
+TVM_REGISTER_OP("relax.nn.dropout")
+    .set_attrs_type<DropoutAttrs>()
+    .set_num_inputs(1)
+    .add_argument("data", "Tensor", "Input to which dropout will be applied.")
+    .set_attr<FInferStructInfo>("FInferStructInfo", InferStructInfoDropout);
+
+}  // namespace relax
+}  // namespace tvm
diff --git a/src/relax/op/nn/nn.h b/src/relax/op/nn/nn.h
new file mode 100644
index 0000000000..df2b978fc2
--- /dev/null
+++ b/src/relax/op/nn/nn.h
@@ -0,0 +1,81 @@
+/*
+ * 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 nn.h
+ * \brief The functions to make Relax neural network operator calls.
+ */
+
+#ifndef TVM_RELAX_OP_NN_NN_H_
+#define TVM_RELAX_OP_NN_NN_H_
+
+#include <tvm/relax/attrs/nn.h>
+
+#include "../op_common.h"
+
+namespace tvm {
+namespace relax {
+
+/*!
+ * \brief Quick helper macro to
+ * - expose a make-function interface which construct the call node.
+ * - register op to the registry.
+ * \param OpName The name of operator to register.
+ * \param OpRegName The identifier of the operator in the registry.
+ * \param RequireFloatDtype A boolean indicating if the input is required to 
have float dtype.
+ */
+#define RELAX_REGISTER_UNARY_NN_OP_AND_IMPL(OpName, OpRegName, 
RequireFloatDtype) \
+  RELAX_REGISTER_UNARY_OP(OpRegName).set_attr<FInferStructInfo>(               
   \
+      "FInferStructInfo", InferStructInfoUnaryArith<RequireFloatDtype>);       
   \
+  RELAX_UNARY_OP_INTERFACE(OpName, OpRegName);
+
+/*! \brief Rectified linear unit. */
+Expr relu(Expr data);
+
+/*! \brief Gaussian Error Linear Units function. */
+Expr gelu(Expr data);
+
+/*! \brief Sigmoid Linear Unit function. */
+Expr silu(Expr data);
+
+/*! \brief Softmax function. */
+Expr softmax(Expr data, int axis);
+
+/*! \brief Compute batch normalization. */
+Expr batch_norm(Expr data, Expr gamma, Expr beta, Expr moving_mean, Expr 
moving_var,  //
+                int axis, double epsilon, bool center, bool scale);
+
+/*! \brief Compute layer normalization. */
+Expr layer_norm(Expr data, Expr gamma, Expr beta, Array<Integer> axes, double 
epsilon, bool center,
+                bool scale);
+
+/*!
+ * \brief Applies the dropout operation to the input tensor.
+ * \param data The input data to the operator.
+ * \param rate The probability for an element to be reset to 0.
+ * \return A Tuple of two tensors.
+ * The first one is the original tensor and the second one is a
+ * mask tensor (1.0 where element not dropped, 0.0 where dropped)
+ */
+Expr dropout(Expr data, double rate);
+
+}  // namespace relax
+}  // namespace tvm
+
+#endif  // TVM_RELAX_OP_NN_NN_H_
diff --git a/src/relax/op/nn/pooling.cc b/src/relax/op/nn/pooling.cc
new file mode 100644
index 0000000000..a4c1e6b17d
--- /dev/null
+++ b/src/relax/op/nn/pooling.cc
@@ -0,0 +1,184 @@
+/*
+ * 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.
+ */
+
+#include "pooling.h"
+
+#include <utility>
+#include <vector>
+
+namespace tvm {
+namespace relax {
+
+/* relax.nn.max_pool2d */
+TVM_REGISTER_NODE_TYPE(MaxPool2DAttrs);
+
+Expr max_pool2d(Expr data, Array<IntImm> pool_size, Array<IntImm> strides, 
Array<IntImm> padding,
+                Array<IntImm> dilation, bool ceil_mode, String layout,
+                Optional<String> out_layout) {
+  padding = GetCompletePadding2D(std::move(padding));
+  if (pool_size.size() == 1) {
+    pool_size.push_back(pool_size[0]);
+  }
+  if (strides.size() == 1) {
+    strides.push_back(strides[0]);
+  }
+  if (dilation.size() == 1) {
+    dilation.push_back(dilation[0]);
+  }
+
+  CHECK_EQ(pool_size.size(), 2)
+      << "The input pool_size length is expected to be 2. However, the given 
pool_size is "
+      << pool_size;
+  CHECK_EQ(strides.size(), 2)
+      << "The input strides length is expected to be 2. However, the given 
strides is " << strides;
+  CHECK_EQ(dilation.size(), 2)
+      << "The input dilation length is expected to be 2. However, the given 
dilation is "
+      << dilation;
+
+  auto attrs = make_object<MaxPool2DAttrs>();
+  attrs->pool_size = std::move(pool_size);
+  attrs->strides = ConvertIntImmToInt64(strides);
+  attrs->padding = ConvertIntImmToInt64(padding);
+  attrs->dilation = ConvertIntImmToInt64(dilation);
+  attrs->ceil_mode = ceil_mode;
+  attrs->layout = layout;
+  attrs->out_layout = out_layout.value_or(layout);
+  static const Op& op = Op::Get("relax.nn.max_pool2d");
+  return Call(op, {std::move(data)}, Attrs(attrs), {});
+}
+
+TVM_REGISTER_GLOBAL("relax.op.nn.max_pool2d").set_body_typed(max_pool2d);
+
+StructInfo InferStructInfoMaxPool2D(const Call& call, const BlockBuilder& ctx) 
{
+  TensorStructInfo data_sinfo = GetUnaryInputTensorStructInfo(call, ctx);
+
+  const auto* attrs = call->attrs.as<MaxPool2DAttrs>();
+  auto [data_layout, data2NCHW] = CheckTensorLayout(call, ctx, attrs->layout,  
//
+                                                    /*tgt_layout=*/"NCHW",     
//
+                                                    /*tensor_name=*/"data");
+  auto [out_layout, out2NCHW] = CheckTensorLayout(call, ctx, 
attrs->out_layout,  //
+                                                  /*tgt_layout=*/"NCHW",       
  //
+                                                  /*tensor_name=*/"output");
+
+  Optional<ShapeExpr> data_shape =
+      CheckNdimPerLayoutAndGetShape(call, ctx, data_sinfo, data_layout);
+  if (!data_shape.defined()) {
+    return TensorStructInfo(data_sinfo->dtype, out_layout.ndim());
+  }
+
+  Array<PrimExpr> data_NCHW_shape = 
data2NCHW.ForwardShape(data_shape.value()->values);
+
+  PrimExpr input_h = data_NCHW_shape[2];
+  PrimExpr input_w = data_NCHW_shape[3];
+  PrimExpr kernel_h = attrs->pool_size[0];
+  PrimExpr kernel_w = attrs->pool_size[1];
+  PrimExpr padding_h = attrs->padding[0] + attrs->padding[2];
+  PrimExpr padding_w = attrs->padding[1] + attrs->padding[3];
+
+  arith::Analyzer* analyzer = ctx->GetAnalyzer();
+  std::vector<PrimExpr> out_NCHW_shape;
+  out_NCHW_shape.resize(4);
+  out_NCHW_shape[0] = data_NCHW_shape[0];
+  out_NCHW_shape[1] = data_NCHW_shape[1];
+
+  PrimExpr numerator_h = input_h + padding_h - attrs->dilation[0] * (kernel_h 
- 1) - 1;
+  PrimExpr numerator_w = input_w + padding_w - attrs->dilation[1] * (kernel_w 
- 1) - 1;
+  if (attrs->ceil_mode) {
+    numerator_h += attrs->strides[0] - 1;
+    numerator_w += attrs->strides[1] - 1;
+  }
+  out_NCHW_shape[2] = analyzer->Simplify(floordiv(numerator_h, 
attrs->strides[0]) + 1);
+  out_NCHW_shape[3] = analyzer->Simplify(floordiv(numerator_w, 
attrs->strides[1]) + 1);
+
+  Array<PrimExpr> out_shape = out2NCHW.BackwardShape(out_NCHW_shape);
+  return TensorStructInfo(ShapeExpr(out_shape), data_sinfo->dtype);
+}
+
+TVM_REGISTER_OP("relax.nn.max_pool2d")
+    .set_num_inputs(1)
+    .add_argument("data", "Tensor", "The input tensor")
+    .set_attrs_type<MaxPool2DAttrs>()
+    .set_attr<FInferStructInfo>("FInferStructInfo", InferStructInfoMaxPool2D);
+
+/* relax.nn.adaptive_avg_pool2d */
+TVM_REGISTER_NODE_TYPE(AdaptivePool2DAttrs);
+
+Expr adaptive_avg_pool2d(Expr data, Optional<Array<IntImm>> output_size, 
String layout,
+                         Optional<String> out_layout) {
+  ObjectPtr<AdaptivePool2DAttrs> attrs = make_object<AdaptivePool2DAttrs>();
+  attrs->layout = layout;
+  attrs->out_layout = out_layout.value_or(layout);
+  if (output_size.defined()) {
+    Array<IntImm> _output_size = output_size.value();
+    if (_output_size.size() == 1) {
+      _output_size.push_back(_output_size[0]);
+    }
+    CHECK_EQ(_output_size.size(), 2)
+        << "The output_size length is expected to be 2. However, the given 
output_size is "
+        << _output_size;
+    attrs->output_size = std::move(_output_size);
+  }
+
+  static const Op& op = Op::Get("relax.nn.adaptive_avg_pool2d");
+  return Call(op, {std::move(data)}, Attrs(attrs), {});
+}
+
+TVM_REGISTER_GLOBAL("relax.op.nn.adaptive_avg_pool2d").set_body_typed(adaptive_avg_pool2d);
+
+StructInfo InferStructInfoAdaptiveAvgPool2D(const Call& call, const 
BlockBuilder& ctx) {
+  TensorStructInfo data_sinfo = GetUnaryInputTensorStructInfo(call, ctx);
+
+  const auto* attrs = call->attrs.as<AdaptivePool2DAttrs>();
+  auto [data_layout, data2NCHW] = CheckTensorLayout(call, ctx, attrs->layout,  
//
+                                                    /*tgt_layout=*/"NCHW",     
//
+                                                    /*tensor_name=*/"data");
+  auto [out_layout, out2NCHW] = CheckTensorLayout(call, ctx, 
attrs->out_layout,  //
+                                                  /*tgt_layout=*/"NCHW",       
  //
+                                                  /*tensor_name=*/"output");
+
+  Optional<ShapeExpr> data_shape =
+      CheckNdimPerLayoutAndGetShape(call, ctx, data_sinfo, data_layout);
+  if (!data_shape.defined()) {
+    if (data_sinfo->shape.defined() && attrs->out_layout == attrs->layout &&
+        !attrs->output_size.defined()) {
+      return data_sinfo;
+    } else {
+      return TensorStructInfo(data_sinfo->dtype, out_layout.ndim());
+    }
+  }
+
+  Array<PrimExpr> data_NCHW_shape = 
data2NCHW.ForwardShape(data_shape.value()->values);
+  Array<PrimExpr> out_NCHW_shape(data_NCHW_shape);
+  if (attrs->output_size.defined()) {
+    out_NCHW_shape.Set(2, attrs->output_size.value()[0]);
+    out_NCHW_shape.Set(3, attrs->output_size.value()[1]);
+  }
+
+  Array<PrimExpr> out_shape = out2NCHW.BackwardShape(out_NCHW_shape);
+  return TensorStructInfo(ShapeExpr(out_shape), data_sinfo->dtype);
+}
+
+TVM_REGISTER_OP("relax.nn.adaptive_avg_pool2d")
+    .set_attrs_type<AdaptivePool2DAttrs>()
+    .set_num_inputs(1)
+    .add_argument("data", "Tensor", "The input tensor")
+    .set_attr<FInferStructInfo>("FInferStructInfo", 
InferStructInfoAdaptiveAvgPool2D);
+
+}  // namespace relax
+}  // namespace tvm
diff --git a/src/relax/op/nn/pooling.h b/src/relax/op/nn/pooling.h
new file mode 100644
index 0000000000..3c1792d21f
--- /dev/null
+++ b/src/relax/op/nn/pooling.h
@@ -0,0 +1,46 @@
+/*
+ * 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 pooling.h
+ * \brief The functions to make Relax neural network pooling operator calls.
+ */
+
+#ifndef TVM_RELAX_OP_NN_POOLING_H_
+#define TVM_RELAX_OP_NN_POOLING_H_
+
+#include <tvm/relax/attrs/nn.h>
+
+#include "../op_common.h"
+
+namespace tvm {
+namespace relax {
+
+/*! \brief 2D maximum pooling operator. */
+Expr max_pool2d(Expr data, Array<IntImm> pool_size, Array<IntImm> strides, 
Array<IntImm> padding,
+                Array<IntImm> dilation, bool ceil_mode, String layout, 
Optional<String> out_layout);
+
+/*! \brief 2D adaptive average pooling operator. */
+Expr adaptive_avg_pool2d(Expr data, Optional<Array<IntImm>> output_size, 
String layout,
+                         Optional<String> out_layout);
+
+}  // namespace relax
+}  // namespace tvm
+
+#endif  // TVM_RELAX_OP_NN_POOLING_H_
diff --git a/tests/python/relax/test_op_nn.py b/tests/python/relax/test_op_nn.py
new file mode 100644
index 0000000000..d047448309
--- /dev/null
+++ b/tests/python/relax/test_op_nn.py
@@ -0,0 +1,929 @@
+# 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.
+import pytest
+import tvm
+import tvm.testing
+from tvm import relax, tir
+from tvm import TVMError
+from tvm.ir import Op
+from tvm.script import relax as R
+
+
+def test_op_correctness():
+    x = relax.Var("x", R.Tensor((2, 3), "float32"))
+    assert relax.op.nn.relu(x).op == Op.get("relax.nn.relu")
+    assert relax.op.nn.gelu(x).op == Op.get("relax.nn.gelu")
+    assert relax.op.nn.silu(x).op == Op.get("relax.nn.silu")
+    assert relax.op.nn.softmax(x).op == Op.get("relax.nn.softmax")
+    assert relax.op.nn.dropout(x).op == Op.get("relax.nn.dropout")
+
+    x = relax.Var("x", R.Tensor((2, 3, 32, 32), "float32"))
+    gamma = relax.Var("gamma", R.Tensor((3,), "float32"))
+    beta = relax.Var("beta", R.Tensor((3,), "float32"))
+    moving_mean = relax.Var("moving_mean", R.Tensor((3,), "float32"))
+    moving_var = relax.Var("moving_var", R.Tensor((3,), "float32"))
+    assert relax.op.nn.batch_norm(x, gamma, beta, moving_mean, moving_var, 
axis=1).op == Op.get(
+        "relax.nn.batch_norm"
+    )
+    assert relax.op.nn.layer_norm(x, gamma, beta, axes=1).op == 
Op.get("relax.nn.layer_norm")
+
+
+def _check_inference(bb: relax.BlockBuilder, call: relax.Call, expected_sinfo: 
relax.StructInfo):
+    ret = bb.normalize(call)
+    tvm.ir.assert_structural_equal(ret.struct_info, expected_sinfo)
+
+
+def test_linear_unit_infer_struct_info():
+    bb = relax.BlockBuilder()
+    x0 = relax.Var("x", R.Tensor((2, 3), "float32"))
+    x1 = relax.Var("x", R.Tensor("float32", ndim=3))
+    x2 = relax.Var("x", R.Tensor("float32", ndim=-1))
+    x3 = relax.Var("x", R.Tensor((2, 3)))
+    x4 = relax.Var("x", R.Tensor())
+
+    _check_inference(bb, relax.op.nn.relu(x0), relax.TensorStructInfo((2, 3), 
"float32"))
+    _check_inference(bb, relax.op.nn.silu(x1), 
relax.TensorStructInfo(dtype="float32", ndim=3))
+    _check_inference(bb, relax.op.nn.gelu(x2), 
relax.TensorStructInfo(dtype="float32"))
+    _check_inference(bb, relax.op.nn.relu(x3), relax.TensorStructInfo((2, 3), 
dtype=""))
+    _check_inference(bb, relax.op.nn.gelu(x4), 
relax.TensorStructInfo(dtype=""))
+
+
+def test_linear_unit_infer_struct_info_shape_symbolic():
+    bb = relax.BlockBuilder()
+    m = tir.Var("m", "int64")
+    n = tir.Var("n", "int64")
+    x0 = relax.Var("x", R.Tensor((m, n), "float32"))
+    x1 = relax.Var("x", R.Tensor((4, n), "float32"))
+
+    _check_inference(bb, relax.op.nn.silu(x0), relax.TensorStructInfo((m, n), 
"float32"))
+    _check_inference(bb, relax.op.nn.relu(x1), relax.TensorStructInfo((4, n), 
"float32"))
+
+
+def test_linear_unit_infer_struct_info_shape_var():
+    bb = relax.BlockBuilder()
+    s0 = relax.Var("s", relax.ShapeStructInfo(ndim=2))
+    s1 = relax.Var("s", relax.ShapeStructInfo())
+    x0 = relax.Var("x", relax.TensorStructInfo(s0, "float32"))
+    x1 = relax.Var("x", relax.TensorStructInfo(s1, "float32"))
+
+    _check_inference(bb, relax.op.nn.gelu(x0), relax.TensorStructInfo(s0, 
"float32"))
+    _check_inference(bb, relax.op.nn.relu(x1), relax.TensorStructInfo(s1, 
"float32"))
+
+
+def test_linear_unit_infer_struct_info_more_input_dtype():
+    bb = relax.BlockBuilder()
+    x0 = relax.Var("x", R.Tensor((2, 3), "float64"))
+    x1 = relax.Var("x", R.Tensor((2, 3), "int8"))
+    x2 = relax.Var("x", R.Tensor((2, 3), "int64"))
+
+    _check_inference(bb, relax.op.nn.relu(x0), relax.TensorStructInfo((2, 3), 
"float64"))
+    _check_inference(bb, relax.op.nn.relu(x1), relax.TensorStructInfo((2, 3), 
"int8"))
+    _check_inference(bb, relax.op.nn.relu(x2), relax.TensorStructInfo((2, 3), 
"int64"))
+
+
+def test_linear_unit_infer_struct_info_invalid_input_dtype():
+    bb = relax.BlockBuilder()
+    x0 = relax.Var("x", R.Tensor((2, 3), "int8"))
+    x1 = relax.Var("x", R.Tensor((2, 3), "int64"))
+
+    with pytest.raises(TVMError):
+        bb.normalize(relax.op.nn.gelu(x0))
+    with pytest.raises(TVMError):
+        bb.normalize(relax.op.nn.silu(x1))
+
+
+def test_linear_unit_infer_struct_info_wrong_input_type():
+    bb = relax.BlockBuilder()
+    x0 = relax.Var("x", relax.ShapeStructInfo((2, 3)))
+    x1 = relax.Var("x", relax.FuncStructInfo([], R.Tensor((2, 3), "float32")))
+
+    with pytest.raises(TVMError):
+        bb.normalize(relax.op.nn.gelu(x0))
+    with pytest.raises(TVMError):
+        bb.normalize(relax.op.nn.silu(x1))
+
+
+def test_softmax_infer_struct_info():
+    bb = relax.BlockBuilder()
+    x0 = relax.Var("x", R.Tensor((2, 3), "float32"))
+    x1 = relax.Var("x", R.Tensor("float32", ndim=3))
+    x2 = relax.Var("x", R.Tensor("float32", ndim=-1))
+    x3 = relax.Var("x", R.Tensor((2, 3)))
+    x4 = relax.Var("x", R.Tensor())
+
+    _check_inference(bb, relax.op.nn.softmax(x0), relax.TensorStructInfo((2, 
3), "float32"))
+    _check_inference(
+        bb, relax.op.nn.softmax(x1, axis=0), 
relax.TensorStructInfo(dtype="float32", ndim=3)
+    )
+    _check_inference(bb, relax.op.nn.softmax(x2, axis=1), 
relax.TensorStructInfo(dtype="float32"))
+    _check_inference(bb, relax.op.nn.softmax(x3, axis=-1), 
relax.TensorStructInfo((2, 3), dtype=""))
+    _check_inference(bb, relax.op.nn.softmax(x4, axis=-2), 
relax.TensorStructInfo(dtype=""))
+
+
+def test_softmax_infer_struct_info_shape_symbolic():
+    bb = relax.BlockBuilder()
+    m = tir.Var("m", "int64")
+    n = tir.Var("n", "int64")
+    x0 = relax.Var("x", R.Tensor((m, n), "float32"))
+    x1 = relax.Var("x", R.Tensor((4, n), "float32"))
+
+    _check_inference(bb, relax.op.nn.softmax(x0), relax.TensorStructInfo((m, 
n), "float32"))
+    _check_inference(bb, relax.op.nn.softmax(x1, axis=0), 
relax.TensorStructInfo((4, n), "float32"))
+
+
+def test_softmax_infer_struct_info_shape_var():
+    bb = relax.BlockBuilder()
+    s0 = relax.Var("s", relax.ShapeStructInfo(ndim=2))
+    s1 = relax.Var("s", relax.ShapeStructInfo())
+    x0 = relax.Var("x", relax.TensorStructInfo(s0, "float32"))
+    x1 = relax.Var("x", relax.TensorStructInfo(s1, "float32"))
+
+    _check_inference(bb, relax.op.nn.softmax(x0), relax.TensorStructInfo(s0, 
"float32"))
+    _check_inference(bb, relax.op.nn.softmax(x1), relax.TensorStructInfo(s1, 
"float32"))
+
+
+def test_softmax_infer_struct_info_more_input_dtype():
+    bb = relax.BlockBuilder()
+    x0 = relax.Var("x", R.Tensor((2, 3), "float16"))
+    x1 = relax.Var("x", R.Tensor((2, 3), "float64"))
+
+    _check_inference(bb, relax.op.nn.softmax(x0), relax.TensorStructInfo((2, 
3), "float16"))
+    _check_inference(bb, relax.op.nn.softmax(x1), relax.TensorStructInfo((2, 
3), "float64"))
+
+
+def test_softmax_infer_struct_info_invalid_input_dtype():
+    bb = relax.BlockBuilder()
+    x0 = relax.Var("x", R.Tensor((2, 3), "int8"))
+    x1 = relax.Var("x", R.Tensor((2, 3), "int64"))
+
+    with pytest.raises(TVMError):
+        bb.normalize(relax.op.nn.softmax(x0))
+    with pytest.raises(TVMError):
+        bb.normalize(relax.op.nn.softmax(x1))
+
+
+def test_softmax_infer_struct_info_axis_out_of_range():
+    bb = relax.BlockBuilder()
+    x = relax.Var("x", R.Tensor((2, 3, 4), "float32"))
+    with pytest.raises(TVMError):
+        bb.normalize(relax.op.nn.softmax(x, axis=3))
+    with pytest.raises(TVMError):
+        bb.normalize(relax.op.nn.softmax(x, axis=-4))
+
+
+def test_softmax_wrong_with_multiple_axes():
+    x = relax.Var("x", R.Tensor((2, 3, 4), "float32"))
+    with pytest.raises(TVMError):
+        relax.op.nn.softmax(x, axis=[1, 2])
+    with pytest.raises(TVMError):
+        relax.op.nn.softmax(x, axis=[-1, -2, -3])
+
+
+def test_softmax_infer_struct_info_wrong_input_type():
+    bb = relax.BlockBuilder()
+    x0 = relax.Var("x", relax.ShapeStructInfo((2, 3)))
+    x1 = relax.Var("x", relax.FuncStructInfo([], R.Tensor((2, 3), "float32")))
+
+    with pytest.raises(TVMError):
+        bb.normalize(relax.op.nn.softmax(x0))
+    with pytest.raises(TVMError):
+        bb.normalize(relax.op.nn.softmax(x1))
+
+
+def test_batch_norm_infer_struct_info():
+    bb = relax.BlockBuilder()
+    x0 = relax.Var("x", R.Tensor((2, 3, 28, 28), "float32"))
+    x1 = relax.Var("x", R.Tensor("float32", ndim=4))
+    x2 = relax.Var("x", R.Tensor("float32"))
+    x3 = relax.Var("x", R.Tensor(ndim=4))
+    x4 = relax.Var("x", R.Tensor())
+    gamma0 = relax.Var("gamma", R.Tensor((3,), "float32"))
+    gamma1 = relax.Var("gamma", R.Tensor("float32", ndim=1))
+    gamma2 = relax.Var("gamma", R.Tensor(ndim=1))
+    beta0 = relax.Var("beta", R.Tensor((3,), "float32"))
+    beta1 = relax.Var("beta", R.Tensor((3,)))
+    moving_mean0 = relax.Var("moving_mean", R.Tensor((3,), "float32"))
+    moving_mean1 = relax.Var("moving_mean", R.Tensor((3,)))
+    moving_var0 = relax.Var("moving_var", R.Tensor((3,), "float32"))
+    moving_var1 = relax.Var("moving_var", R.Tensor("float32", ndim=1))
+    moving_var2 = relax.Var("moving_var", R.Tensor(ndim=1))
+
+    _check_inference(
+        bb,
+        relax.op.nn.batch_norm(x0, gamma0, beta0, moving_mean0, moving_var0, 
axis=1),
+        relax.TupleStructInfo(
+            [
+                relax.TensorStructInfo((2, 3, 28, 28), "float32"),
+                relax.TensorStructInfo((3,), "float32"),
+                relax.TensorStructInfo((3,), "float32"),
+            ]
+        ),
+    )
+    _check_inference(
+        bb,
+        relax.op.nn.batch_norm(x0, gamma0, beta0, moving_mean0, moving_var0, 
axis=-3),
+        relax.TupleStructInfo(
+            [
+                relax.TensorStructInfo((2, 3, 28, 28), "float32"),
+                relax.TensorStructInfo((3,), "float32"),
+                relax.TensorStructInfo((3,), "float32"),
+            ]
+        ),
+    )
+    _check_inference(
+        bb,
+        relax.op.nn.batch_norm(x1, gamma0, beta0, moving_mean0, moving_var0, 
axis=1),
+        relax.TupleStructInfo(
+            [
+                relax.TensorStructInfo(dtype="float32", ndim=4),
+                relax.TensorStructInfo((3,), "float32"),
+                relax.TensorStructInfo((3,), "float32"),
+            ]
+        ),
+    )
+    _check_inference(
+        bb,
+        relax.op.nn.batch_norm(x0, gamma1, beta0, moving_mean0, moving_var0, 
axis=1),
+        relax.TupleStructInfo(
+            [
+                relax.TensorStructInfo((2, 3, 28, 28), "float32"),
+                relax.TensorStructInfo((3,), "float32"),
+                relax.TensorStructInfo((3,), "float32"),
+            ]
+        ),
+    )
+    _check_inference(
+        bb,
+        relax.op.nn.batch_norm(x0, gamma0, beta0, moving_mean0, moving_var1, 
axis=1),
+        relax.TupleStructInfo(
+            [
+                relax.TensorStructInfo((2, 3, 28, 28), "float32"),
+                relax.TensorStructInfo((3,), "float32"),
+                relax.TensorStructInfo(dtype="float32", ndim=1),
+            ]
+        ),
+    )
+    _check_inference(
+        bb,
+        relax.op.nn.batch_norm(x1, gamma1, beta0, moving_mean0, moving_var1, 
axis=1),
+        relax.TupleStructInfo(
+            [
+                relax.TensorStructInfo(dtype="float32", ndim=4),
+                relax.TensorStructInfo((3,), "float32"),
+                relax.TensorStructInfo(dtype="float32", ndim=1),
+            ]
+        ),
+    )
+    _check_inference(
+        bb,
+        relax.op.nn.batch_norm(x2, gamma1, beta0, moving_mean0, moving_var1, 
axis=1),
+        relax.TupleStructInfo(
+            [
+                relax.TensorStructInfo(dtype="float32"),
+                relax.TensorStructInfo((3,), "float32"),
+                relax.TensorStructInfo(dtype="float32", ndim=1),
+            ]
+        ),
+    )
+    _check_inference(
+        bb,
+        relax.op.nn.batch_norm(x3, gamma2, beta1, moving_mean1, moving_var2, 
axis=1),
+        relax.TupleStructInfo(
+            [
+                relax.TensorStructInfo(ndim=4, dtype=""),
+                relax.TensorStructInfo((3,), dtype=""),
+                relax.TensorStructInfo(dtype="", ndim=1),
+            ]
+        ),
+    )
+    _check_inference(
+        bb,
+        relax.op.nn.batch_norm(x4, gamma2, beta1, moving_mean1, moving_var2, 
axis=1),
+        relax.TupleStructInfo(
+            [
+                relax.TensorStructInfo(dtype=""),
+                relax.TensorStructInfo((3,), dtype=""),
+                relax.TensorStructInfo(dtype="", ndim=1),
+            ]
+        ),
+    )
+
+
+def test_batch_norm_infer_struct_info_shape_symbolic():
+    bb = relax.BlockBuilder()
+    n = tir.Var("n", "int64")
+    c0 = tir.Var("c", "int64")
+    c1 = tir.Var("c", "int64")
+    h = tir.Var("h", "int64")
+    w = tir.Var("w", "int64")
+    x0 = relax.Var("x", R.Tensor((n, c0, h, w), "float32"))
+    x1 = relax.Var("x", R.Tensor((n, c1, h, w), "float32"))
+    x2 = relax.Var("x", R.Tensor("float32", ndim=4))
+    gamma0 = relax.Var("gamma", R.Tensor((c0,), "float32"))
+    gamma1 = relax.Var("gamma", R.Tensor((c1,), "float32"))
+    gamma2 = relax.Var("gamma", R.Tensor("float32", ndim=1))
+    beta = relax.Var("beta", R.Tensor((c0,), "float32"))
+    moving_mean = relax.Var("moving_mean", R.Tensor((c0,), "float32"))
+    moving_var0 = relax.Var("moving_var", R.Tensor((c0,), "float32"))
+    moving_var1 = relax.Var("moving_var", R.Tensor((c1,), "float32"))
+    moving_var2 = relax.Var("moving_var", R.Tensor("float32", ndim=1))
+
+    _check_inference(
+        bb,
+        relax.op.nn.batch_norm(x0, gamma0, beta, moving_mean, moving_var0, 
axis=1),
+        relax.TupleStructInfo(
+            [
+                relax.TensorStructInfo((n, c0, h, w), "float32"),
+                relax.TensorStructInfo((c0,), "float32"),
+                relax.TensorStructInfo((c0,), "float32"),
+            ]
+        ),
+    )
+    _check_inference(
+        bb,
+        relax.op.nn.batch_norm(x1, gamma0, beta, moving_mean, moving_var0, 
axis=1),
+        relax.TupleStructInfo(
+            [
+                relax.TensorStructInfo(dtype="float32", ndim=4),
+                relax.TensorStructInfo(dtype="float32", ndim=1),
+                relax.TensorStructInfo(dtype="float32", ndim=1),
+            ]
+        ),
+    )
+    _check_inference(
+        bb,
+        relax.op.nn.batch_norm(x2, gamma0, beta, moving_mean, moving_var0, 
axis=1),
+        relax.TupleStructInfo(
+            [
+                relax.TensorStructInfo(dtype="float32", ndim=4),
+                relax.TensorStructInfo((c0,), "float32"),
+                relax.TensorStructInfo((c0,), "float32"),
+            ]
+        ),
+    )
+    _check_inference(
+        bb,
+        relax.op.nn.batch_norm(x0, gamma1, beta, moving_mean, moving_var0, 
axis=1),
+        relax.TupleStructInfo(
+            [
+                relax.TensorStructInfo(dtype="float32", ndim=4),
+                relax.TensorStructInfo(dtype="float32", ndim=1),
+                relax.TensorStructInfo(dtype="float32", ndim=1),
+            ]
+        ),
+    )
+    _check_inference(
+        bb,
+        relax.op.nn.batch_norm(x0, gamma0, beta, moving_mean, moving_var1, 
axis=1),
+        relax.TupleStructInfo(
+            [
+                relax.TensorStructInfo(dtype="float32", ndim=4),
+                relax.TensorStructInfo(dtype="float32", ndim=1),
+                relax.TensorStructInfo(dtype="float32", ndim=1),
+            ]
+        ),
+    )
+    _check_inference(
+        bb,
+        relax.op.nn.batch_norm(x0, gamma2, beta, moving_mean, moving_var0, 
axis=1),
+        relax.TupleStructInfo(
+            [
+                relax.TensorStructInfo((n, c0, h, w), "float32"),
+                relax.TensorStructInfo((c0,), "float32"),
+                relax.TensorStructInfo((c0,), "float32"),
+            ]
+        ),
+    )
+    _check_inference(
+        bb,
+        relax.op.nn.batch_norm(x0, gamma0, beta, moving_mean, moving_var2, 
axis=1),
+        relax.TupleStructInfo(
+            [
+                relax.TensorStructInfo((n, c0, h, w), "float32"),
+                relax.TensorStructInfo((c0,), "float32"),
+                relax.TensorStructInfo(dtype="float32", ndim=1),
+            ]
+        ),
+    )
+
+
+def test_batch_norm_infer_struct_info_shape_var():
+    bb = relax.BlockBuilder()
+    s0 = relax.Var("s0", relax.ShapeStructInfo(ndim=4))
+    s1 = relax.Var("s1", relax.ShapeStructInfo())
+    s2 = relax.Var("s2", relax.ShapeStructInfo(ndim=1))
+    s3 = relax.Var("s3", relax.ShapeStructInfo(ndim=1))
+    x0 = relax.Var("x", relax.TensorStructInfo(s0, "float32"))
+    x1 = relax.Var("x", relax.TensorStructInfo(s1, "float32"))
+    gamma = relax.Var("gamma", relax.TensorStructInfo(s2, "float32"))
+    beta = relax.Var("beta", relax.TensorStructInfo(s3, "float32"))
+    moving_mean = relax.Var("moving_mean", relax.TensorStructInfo(s2, 
"float32"))
+    moving_var = relax.Var("moving_var", relax.TensorStructInfo(s3, "float32"))
+
+    _check_inference(
+        bb,
+        relax.op.nn.batch_norm(x0, gamma, beta, moving_mean, moving_var, 
axis=1),
+        relax.TupleStructInfo(
+            [
+                relax.TensorStructInfo(s0, "float32"),
+                relax.TensorStructInfo(s2, "float32"),
+                relax.TensorStructInfo(s3, "float32"),
+            ]
+        ),
+    )
+    _check_inference(
+        bb,
+        relax.op.nn.batch_norm(x1, gamma, beta, moving_mean, moving_var, 
axis=1),
+        relax.TupleStructInfo(
+            [
+                relax.TensorStructInfo(s1, "float32"),
+                relax.TensorStructInfo(s2, "float32"),
+                relax.TensorStructInfo(s3, "float32"),
+            ]
+        ),
+    )
+
+
+def test_batch_norm_infer_struct_info_more_input_dtype():
+    bb = relax.BlockBuilder()
+    x = relax.Var("x", R.Tensor((2, 3, 28, 28), "float16"))
+    gamma = relax.Var("gamma", R.Tensor((3,), "float16"))
+    beta = relax.Var("beta", R.Tensor((3,), "float16"))
+    moving_mean = relax.Var("moving_mean", R.Tensor((3,), "float16"))
+    moving_var = relax.Var("moving_var", R.Tensor((3,), "float16"))
+
+    _check_inference(
+        bb,
+        relax.op.nn.batch_norm(x, gamma, beta, moving_mean, moving_var, 
axis=1),
+        relax.TupleStructInfo(
+            [
+                relax.TensorStructInfo((2, 3, 28, 28), "float16"),
+                relax.TensorStructInfo((3,), "float16"),
+                relax.TensorStructInfo((3,), "float16"),
+            ]
+        ),
+    )
+
+
+def test_batch_norm_infer_struct_info_invalid_input_dtype():
+    bb = relax.BlockBuilder()
+    x0 = relax.Var("x", R.Tensor((2, 3, 28, 28), "int8"))
+    gamma0 = relax.Var("gamma", R.Tensor((3,), "int8"))
+    beta0 = relax.Var("beta", R.Tensor((3,), "int8"))
+    moving_mean0 = relax.Var("moving_mean", R.Tensor((3,), "int8"))
+    moving_var0 = relax.Var("moving_var", R.Tensor((3,), "int8"))
+    x1 = relax.Var("x", R.Tensor((2, 3, 28, 28), "int32"))
+    gamma1 = relax.Var("gamma", R.Tensor((3,), "int32"))
+    beta1 = relax.Var("beta", R.Tensor((3,), "int32"))
+    moving_mean1 = relax.Var("moving_mean", R.Tensor((3,), "int32"))
+    moving_var1 = relax.Var("moving_var", R.Tensor((3,), "int32"))
+
+    with pytest.raises(TVMError):
+        bb.normalize(relax.op.nn.batch_norm(x0, gamma0, beta0, moving_mean0, 
moving_var0, axis=1))
+    with pytest.raises(TVMError):
+        bb.normalize(relax.op.nn.batch_norm(x1, gamma1, beta1, moving_mean1, 
moving_var1, axis=1))
+
+
+def test_batch_norm_infer_struct_info_axis_out_of_range():
+    bb = relax.BlockBuilder()
+    x = relax.Var("x", R.Tensor((2, 3, 28, 28), "float32"))
+    gamma = relax.Var("gamma", R.Tensor((3,), "float32"))
+    beta = relax.Var("beta", R.Tensor((3,), "float32"))
+    moving_mean = relax.Var("moving_mean", R.Tensor((3,), "float32"))
+    moving_var = relax.Var("moving_var", R.Tensor((3,), "float32"))
+
+    with pytest.raises(TVMError):
+        bb.normalize(relax.op.nn.batch_norm(x, gamma, beta, moving_mean, 
moving_var, axis=4))
+    with pytest.raises(TVMError):
+        bb.normalize(relax.op.nn.batch_norm(x, gamma, beta, moving_mean, 
moving_var, axis=-5))
+
+
+def test_batch_norm_infer_struct_info_dtype_mismatch():
+    bb = relax.BlockBuilder()
+    x0 = relax.Var("x", R.Tensor((2, 3, 28, 28), "float32"))
+    x1 = relax.Var("x", R.Tensor((2, 3, 28, 28), "int8"))
+    gamma0 = relax.Var("gamma", R.Tensor((3,), "float32"))
+    gamma1 = relax.Var("gamma", R.Tensor((3,)))
+    beta = relax.Var("beta", R.Tensor((3,), "float32"))
+    moving_mean = relax.Var("moving_mean", R.Tensor((3,), "float32"))
+    moving_var0 = relax.Var("moving_var", R.Tensor((3,), "float32"))
+    moving_var1 = relax.Var("moving_var", R.Tensor((3,), "float16"))
+
+    with pytest.raises(TVMError):
+        bb.normalize(relax.op.nn.batch_norm(x1, gamma0, beta, moving_mean, 
moving_var0, axis=1))
+    with pytest.raises(TVMError):
+        bb.normalize(relax.op.nn.batch_norm(x0, gamma1, beta, moving_mean, 
moving_var0, axis=1))
+    with pytest.raises(TVMError):
+        bb.normalize(relax.op.nn.batch_norm(x0, gamma0, beta, moving_mean, 
moving_var1, axis=1))
+
+
+def test_batch_norm_infer_struct_info_ndim_mismatch():
+    bb = relax.BlockBuilder()
+    x = relax.Var("x", R.Tensor((2, 3, 28, 28), "float32"))
+    gamma0 = relax.Var("gamma", R.Tensor((3,), "float32"))
+    gamma1 = relax.Var("gamma", R.Tensor((3, 1), "float32"))
+    beta = relax.Var("beta", R.Tensor((3,), "float32"))
+    moving_mean = relax.Var("moving_mean", R.Tensor((3,), "float32"))
+    moving_var0 = relax.Var("moving_var", R.Tensor((3,), "float32"))
+    moving_var1 = relax.Var("moving_var", R.Tensor((1, 3), "float32"))
+
+    with pytest.raises(TVMError):
+        bb.normalize(relax.op.nn.batch_norm(x, gamma1, beta, moving_mean, 
moving_var0, axis=1))
+    with pytest.raises(TVMError):
+        bb.normalize(relax.op.nn.batch_norm(x, gamma0, beta, moving_mean, 
moving_var1, axis=1))
+
+
+def test_batch_norm_infer_struct_info_shape_mismatch():
+    bb = relax.BlockBuilder()
+    c = tir.Var("c", "int64")
+    x0 = relax.Var("x", R.Tensor((2, 3, 28, 28), "float32"))
+    x1 = relax.Var("x", R.Tensor((2, c, 28, 28), "float32"))
+    gamma0 = relax.Var("gamma", R.Tensor((3,), "float32"))
+    gamma1 = relax.Var("gamma", R.Tensor((4,), "float32"))
+    gamma2 = relax.Var("gamma", R.Tensor((c + 2,), "float32"))
+    beta0 = relax.Var("beta", R.Tensor((3,), "float32"))
+    beta1 = relax.Var("beta", R.Tensor((c,), "float32"))
+    moving_mean0 = relax.Var("moving_mean", R.Tensor((3,), "float32"))
+    moving_mean1 = relax.Var("moving_mean", R.Tensor((c,), "float32"))
+    moving_var0 = relax.Var("moving_var", R.Tensor((3,), "float32"))
+    moving_var1 = relax.Var("moving_var", R.Tensor((4,), "float32"))
+    moving_var2 = relax.Var("moving_var", R.Tensor((c,), "float32"))
+
+    with pytest.raises(TVMError):
+        bb.normalize(relax.op.nn.batch_norm(x0, gamma1, beta0, moving_mean0, 
moving_var0, axis=1))
+    with pytest.raises(TVMError):
+        bb.normalize(relax.op.nn.batch_norm(x0, gamma0, beta0, moving_mean0, 
moving_var1, axis=1))
+    with pytest.raises(TVMError):
+        bb.normalize(relax.op.nn.batch_norm(x1, gamma2, beta1, moving_mean1, 
moving_var2, axis=1))
+
+
+def test_batch_norm_infer_struct_info_wrong_input_type():
+    bb = relax.BlockBuilder()
+    x0 = relax.Var("x", R.Tensor((2, 3, 28, 28), "float32"))
+    x1 = relax.Var("x", relax.ShapeStructInfo((2, 3, 28, 28)))
+    gamma0 = relax.Var("gamma", R.Tensor((3,), "float32"))
+    gamma1 = relax.Var("gamma", relax.FuncStructInfo([], R.Tensor((3,), 
"float32")))
+    beta = relax.Var("beta", R.Tensor((3,), "float32"))
+    moving_mean = relax.Var("moving_mean", R.Tensor((3,), "float32"))
+    moving_var = relax.Var("moving_var", R.Tensor((3,), "float32"))
+
+    with pytest.raises(TVMError):
+        bb.normalize(relax.op.nn.batch_norm(x1, gamma0, beta, moving_mean, 
moving_var, axis=1))
+    with pytest.raises(TVMError):
+        bb.normalize(relax.op.nn.batch_norm(x0, gamma1, beta, moving_mean, 
moving_var, axis=1))
+
+
+def test_layer_norm_infer_struct_info():
+    bb = relax.BlockBuilder()
+    x0 = relax.Var("x", R.Tensor((2, 3, 4, 5), "float32"))
+    x1 = relax.Var("x", R.Tensor("float32", ndim=4))
+    x2 = relax.Var("x", R.Tensor("float32"))
+    x3 = relax.Var("x", R.Tensor((2, 3, 4, 5)))
+    gamma0 = relax.Var("gamma", R.Tensor((4, 5), "float32"))
+    gamma1 = relax.Var("gamma", R.Tensor("float32", ndim=2))
+    gamma2 = relax.Var("gamma", R.Tensor((4, 5)))
+    beta0 = relax.Var("beta", R.Tensor((4, 5), "float32"))
+    beta1 = relax.Var("beta", R.Tensor((4, 5)))
+
+    _check_inference(
+        bb,
+        relax.op.nn.layer_norm(x0, gamma0, beta0, axes=[-2, -1]),
+        relax.TensorStructInfo((2, 3, 4, 5), "float32"),
+    )
+    _check_inference(
+        bb,
+        relax.op.nn.layer_norm(x0, gamma0, beta0, axes=[-2, 3]),
+        relax.TensorStructInfo((2, 3, 4, 5), "float32"),
+    )
+    _check_inference(
+        bb,
+        relax.op.nn.layer_norm(x1, gamma0, beta0, axes=[-2, -1]),
+        relax.TensorStructInfo(dtype="float32", ndim=4),
+    )
+    _check_inference(
+        bb,
+        relax.op.nn.layer_norm(x2, gamma0, beta0, axes=[-2, -1]),
+        relax.TensorStructInfo(dtype="float32"),
+    )
+    _check_inference(
+        bb,
+        relax.op.nn.layer_norm(x0, gamma1, beta0, axes=[-2, -1]),
+        relax.TensorStructInfo((2, 3, 4, 5), dtype="float32"),
+    )
+    _check_inference(
+        bb,
+        relax.op.nn.layer_norm(x3, gamma2, beta1, axes=[-2, -1]),
+        relax.TensorStructInfo((2, 3, 4, 5), dtype=""),
+    )
+
+
+def test_layer_norm_infer_struct_info_shape_symbolic():
+    bb = relax.BlockBuilder()
+    n = tir.Var("n", "int64")
+    a = tir.Var("a", "int64")
+    b = tir.Var("b", "int64")
+    c0 = tir.Var("c", "int64")
+    c1 = tir.Var("c", "int64")
+    x0 = relax.Var("x", R.Tensor((n, a, b, c0), "float32"))
+    x1 = relax.Var("x", R.Tensor((n, a, b, c1), "float32"))
+    x2 = relax.Var("x", R.Tensor("float32", ndim=4))
+    gamma0 = relax.Var("gamma", R.Tensor((b, c0), "float32"))
+    gamma1 = relax.Var("gamma", R.Tensor((b, c1), "float32"))
+    beta = relax.Var("beta", R.Tensor((b, c0), "float32"))
+
+    _check_inference(
+        bb,
+        relax.op.nn.layer_norm(x0, gamma0, beta, axes=[-2, -1]),
+        relax.TensorStructInfo((n, a, b, c0), "float32"),
+    )
+    _check_inference(
+        bb,
+        relax.op.nn.layer_norm(x1, gamma0, beta, axes=[-2, -1]),
+        relax.TensorStructInfo(dtype="float32", ndim=4),
+    )
+    _check_inference(
+        bb,
+        relax.op.nn.layer_norm(x0, gamma1, beta, axes=[-2, -1]),
+        relax.TensorStructInfo(dtype="float32", ndim=4),
+    )
+    _check_inference(
+        bb,
+        relax.op.nn.layer_norm(x2, gamma0, beta, axes=[-2, -1]),
+        relax.TensorStructInfo(dtype="float32", ndim=4),
+    )
+    _check_inference(
+        bb,
+        relax.op.nn.layer_norm(x2, gamma1, beta, axes=[-2, -1]),
+        relax.TensorStructInfo(dtype="float32", ndim=4),
+    )
+
+
+def test_layer_norm_infer_struct_info_shape_var():
+    bb = relax.BlockBuilder()
+    s0 = relax.Var("s0", relax.ShapeStructInfo(ndim=4))
+    s1 = relax.Var("s1", relax.ShapeStructInfo())
+    s2 = relax.Var("s2", relax.ShapeStructInfo(ndim=2))
+    s3 = relax.Var("s3", relax.ShapeStructInfo(ndim=2))
+    x0 = relax.Var("x", relax.TensorStructInfo(s0, "float32"))
+    x1 = relax.Var("x", relax.TensorStructInfo(s1, "float32"))
+    gamma = relax.Var("gamma", relax.TensorStructInfo(s2, "float32"))
+    beta = relax.Var("beta", relax.TensorStructInfo(s3, "float32"))
+
+    _check_inference(
+        bb,
+        relax.op.nn.layer_norm(x0, gamma, beta, axes=[2, 3]),
+        relax.TensorStructInfo(s0, "float32"),
+    )
+    _check_inference(
+        bb,
+        relax.op.nn.layer_norm(x1, gamma, beta, axes=[2, 3]),
+        relax.TensorStructInfo(s1, "float32"),
+    )
+
+
+def test_layer_norm_infer_struct_info_more_input_dtype():
+    bb = relax.BlockBuilder()
+    x0 = relax.Var("x", R.Tensor((2, 3, 4, 5), "float16"))
+    gamma0 = relax.Var("gamma", R.Tensor((4, 5), "float16"))
+    beta0 = relax.Var("beta", R.Tensor((4, 5), "float16"))
+    x1 = relax.Var("x", R.Tensor((2, 3, 4, 5), "float64"))
+    gamma1 = relax.Var("gamma", R.Tensor((4, 5), "float64"))
+    beta1 = relax.Var("beta", R.Tensor((4, 5), "float64"))
+
+    _check_inference(
+        bb,
+        relax.op.nn.layer_norm(x0, gamma0, beta0, axes=[-2, -1]),
+        relax.TensorStructInfo((2, 3, 4, 5), "float16"),
+    )
+    _check_inference(
+        bb,
+        relax.op.nn.layer_norm(x1, gamma1, beta1, axes=[-2, -1]),
+        relax.TensorStructInfo((2, 3, 4, 5), "float64"),
+    )
+
+
+def test_layer_norm_infer_struct_info_invalid_input_dtype():
+    bb = relax.BlockBuilder()
+    x0 = relax.Var("x", R.Tensor((2, 3, 4, 5), "int8"))
+    gamma0 = relax.Var("gamma", R.Tensor((4, 5), "int8"))
+    beta0 = relax.Var("beta", R.Tensor((4, 5), "int8"))
+    x1 = relax.Var("x", R.Tensor((2, 3, 4, 5), "int32"))
+    gamma1 = relax.Var("gamma", R.Tensor((4, 5), "int32"))
+    beta1 = relax.Var("beta", R.Tensor((4, 5), "int32"))
+
+    with pytest.raises(TVMError):
+        bb.normalize(relax.op.nn.layer_norm(x0, gamma0, beta0, axes=[-2, -1]))
+    with pytest.raises(TVMError):
+        bb.normalize(relax.op.nn.layer_norm(x1, gamma1, beta1, axes=[-2, -1]))
+
+
+def test_layer_norm_infer_struct_info_axis_out_of_range_and_repetitive():
+    bb = relax.BlockBuilder()
+    x = relax.Var("x", R.Tensor((2, 3, 4, 5), "float32"))
+    gamma = relax.Var("gamma", R.Tensor((4, 5), "float32"))
+    beta = relax.Var("beta", R.Tensor((4, 5), "float32"))
+
+    with pytest.raises(TVMError):
+        bb.normalize(relax.op.nn.layer_norm(x, gamma, beta, axes=[3, 4]))
+    with pytest.raises(TVMError):
+        bb.normalize(relax.op.nn.layer_norm(x, gamma, beta, axes=[3, -1]))
+
+
+def test_layer_norm_infer_struct_info_dtype_mismatch():
+    bb = relax.BlockBuilder()
+    x = relax.Var("x", R.Tensor((2, 3, 4, 5), "float32"))
+    gamma0 = relax.Var("gamma", R.Tensor((4, 5), "float32"))
+    gamma1 = relax.Var("gamma", R.Tensor((4, 5), "int8"))
+    beta0 = relax.Var("beta", R.Tensor((4, 5), "float32"))
+    beta1 = relax.Var("beta", R.Tensor((4, 5)))
+
+    with pytest.raises(TVMError):
+        bb.normalize(relax.op.nn.layer_norm(x, gamma1, beta0, axes=[-2, -1]))
+    with pytest.raises(TVMError):
+        bb.normalize(relax.op.nn.layer_norm(x, gamma0, beta1, axes=[-2, -1]))
+
+
+def test_layer_norm_infer_struct_info_ndim_mismatch():
+    bb = relax.BlockBuilder()
+    x = relax.Var("x", R.Tensor((2, 3, 4, 5), "float32"))
+    gamma0 = relax.Var("gamma", R.Tensor((4, 5), "float32"))
+    gamma1 = relax.Var("gamma", R.Tensor((4,), "float32"))
+    beta0 = relax.Var("beta", R.Tensor((4, 5), "float32"))
+    beta1 = relax.Var("beta", R.Tensor((3, 4, 5), "float32"))
+
+    with pytest.raises(TVMError):
+        bb.normalize(relax.op.nn.layer_norm(x, gamma1, beta0, axes=[-2, -1]))
+    with pytest.raises(TVMError):
+        bb.normalize(relax.op.nn.layer_norm(x, gamma0, beta1, axes=[-2, -1]))
+
+
+def test_layer_norm_infer_struct_info_shape_mismatch():
+    bb = relax.BlockBuilder()
+    c0 = tir.Var("c", "int64")
+    x0 = relax.Var("x", R.Tensor((2, 3, 4, 5), "float32"))
+    x1 = relax.Var("x", R.Tensor((2, 3, 4, c0), "float32"))
+    gamma0 = relax.Var("gamma", R.Tensor((4, 6), "float32"))
+    gamma1 = relax.Var("gamma", R.Tensor((4, c0), "float32"))
+    beta0 = relax.Var("beta", R.Tensor((4, 5), "float32"))
+    beta1 = relax.Var("beta", R.Tensor((4, c0 - 2), "float32"))
+
+    with pytest.raises(TVMError):
+        bb.normalize(relax.op.nn.layer_norm(x0, gamma0, beta0, axes=[-2, -1]))
+    with pytest.raises(TVMError):
+        bb.normalize(relax.op.nn.layer_norm(x1, gamma1, beta1, axes=[-2, -1]))
+
+
+def test_layer_norm_infer_struct_info_wrong_input_type():
+    bb = relax.BlockBuilder()
+    x0 = relax.Var("x", R.Tensor((2, 3, 4, 5), "float32"))
+    x1 = relax.Var("x", relax.ShapeStructInfo((2, 3, 4, 5)))
+    gamma0 = relax.Var("gamma", R.Tensor((4, 5), "float32"))
+    gamma1 = relax.Var("gamma", relax.FuncStructInfo([], R.Tensor((4, 5), 
"float32")))
+    beta = relax.Var("beta", R.Tensor((4, 5), "float32"))
+
+    with pytest.raises(TVMError):
+        bb.normalize(relax.op.nn.layer_norm(x1, gamma0, beta, axes=[-2, -1]))
+    with pytest.raises(TVMError):
+        bb.normalize(relax.op.nn.layer_norm(x0, gamma1, beta, axes=[-2, -1]))
+
+
+def test_dropout_infer_struct_info():
+    bb = relax.BlockBuilder()
+    x0 = relax.Var("x", R.Tensor((2, 3), "float32"))
+    x1 = relax.Var("x", R.Tensor("float32", ndim=3))
+    x2 = relax.Var("x", R.Tensor("float32", ndim=-1))
+    x3 = relax.Var("x", R.Tensor((2, 3)))
+    x4 = relax.Var("x", R.Tensor())
+
+    _check_inference(
+        bb,
+        relax.op.nn.dropout(x0),
+        relax.TupleStructInfo(
+            [relax.TensorStructInfo((2, 3), "float32"), 
relax.TensorStructInfo((2, 3), "float32")]
+        ),
+    )
+    _check_inference(
+        bb,
+        relax.op.nn.dropout(x1),
+        relax.TupleStructInfo(
+            [
+                relax.TensorStructInfo(dtype="float32", ndim=3),
+                relax.TensorStructInfo(dtype="float32", ndim=3),
+            ]
+        ),
+    )
+    _check_inference(
+        bb,
+        relax.op.nn.dropout(x2),
+        relax.TupleStructInfo(
+            [relax.TensorStructInfo(dtype="float32"), 
relax.TensorStructInfo(dtype="float32")]
+        ),
+    )
+    _check_inference(
+        bb,
+        relax.op.nn.dropout(x3),
+        relax.TupleStructInfo(
+            [relax.TensorStructInfo((2, 3), dtype=""), 
relax.TensorStructInfo((2, 3), dtype="")]
+        ),
+    )
+    _check_inference(
+        bb,
+        relax.op.nn.dropout(x4),
+        relax.TupleStructInfo([relax.TensorStructInfo(dtype=""), 
relax.TensorStructInfo(dtype="")]),
+    )
+
+
+def test_dropout_infer_struct_info_shape_symbolic():
+    bb = relax.BlockBuilder()
+    m = tir.Var("m", "int64")
+    n = tir.Var("n", "int64")
+    x = relax.Var("x", R.Tensor((m, n), "float32"))
+
+    _check_inference(
+        bb,
+        relax.op.nn.dropout(x),
+        relax.TupleStructInfo(
+            [relax.TensorStructInfo((m, n), "float32"), 
relax.TensorStructInfo((m, n), "float32")]
+        ),
+    )
+
+
+def test_dropout_infer_struct_info_shape_var():
+    bb = relax.BlockBuilder()
+    s0 = relax.Var("s", relax.ShapeStructInfo(ndim=2))
+    s1 = relax.Var("s", relax.ShapeStructInfo())
+    x0 = relax.Var("x", relax.TensorStructInfo(s0, "float32"))
+    x1 = relax.Var("x", relax.TensorStructInfo(s1, "float32"))
+
+    _check_inference(
+        bb,
+        relax.op.nn.dropout(x0),
+        relax.TupleStructInfo(
+            [relax.TensorStructInfo(s0, "float32"), relax.TensorStructInfo(s0, 
"float32")]
+        ),
+    )
+    _check_inference(
+        bb,
+        relax.op.nn.dropout(x1),
+        relax.TupleStructInfo(
+            [relax.TensorStructInfo(s1, "float32"), relax.TensorStructInfo(s1, 
"float32")]
+        ),
+    )
+
+
+def test_dropout_infer_struct_info_more_input_dtype():
+    bb = relax.BlockBuilder()
+    x0 = relax.Var("x", R.Tensor((2, 3), "float64"))
+    x1 = relax.Var("x", R.Tensor((2, 3), "int8"))
+    x2 = relax.Var("x", R.Tensor((2, 3), "int64"))
+
+    _check_inference(
+        bb,
+        relax.op.nn.dropout(x0),
+        relax.TupleStructInfo(
+            [relax.TensorStructInfo((2, 3), "float64"), 
relax.TensorStructInfo((2, 3), "float64")]
+        ),
+    )
+    _check_inference(
+        bb,
+        relax.op.nn.dropout(x1),
+        relax.TupleStructInfo(
+            [relax.TensorStructInfo((2, 3), "int8"), 
relax.TensorStructInfo((2, 3), "int8")]
+        ),
+    )
+    _check_inference(
+        bb,
+        relax.op.nn.dropout(x2),
+        relax.TupleStructInfo(
+            [relax.TensorStructInfo((2, 3), "int64"), 
relax.TensorStructInfo((2, 3), "int64")]
+        ),
+    )
+
+
+def test_dropout_infer_struct_info_wrong_input_type():
+    bb = relax.BlockBuilder()
+    x0 = relax.Var("x", relax.ShapeStructInfo((2, 3)))
+    x1 = relax.Var("x", relax.FuncStructInfo([], R.Tensor((2, 3), "float32")))
+
+    with pytest.raises(TVMError):
+        bb.normalize(relax.op.nn.dropout(x0))
+    with pytest.raises(TVMError):
+        bb.normalize(relax.op.nn.dropout(x1))
+
+
+if __name__ == "__main__":
+    tvm.testing.main()
diff --git a/tests/python/relax/test_op_nn_convolution.py 
b/tests/python/relax/test_op_nn_convolution.py
new file mode 100644
index 0000000000..6533d43420
--- /dev/null
+++ b/tests/python/relax/test_op_nn_convolution.py
@@ -0,0 +1,429 @@
+# 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.
+import pytest
+import tvm
+import tvm.testing
+from tvm import relax, tir
+from tvm import TVMError
+from tvm.ir import Op
+from tvm.script import relax as R
+
+
+def test_op_correctness():
+    x = relax.Var("x", R.Tensor((2, 3, 28, 28), "float32"))
+    w = relax.Var("w", R.Tensor((4, 3, 3, 3), "float32"))
+    assert relax.op.nn.conv2d(x, w).op == Op.get("relax.nn.conv2d")
+
+
+def _check_inference(bb: relax.BlockBuilder, call: relax.Call, expected_sinfo: 
relax.StructInfo):
+    ret = bb.normalize(call)
+    tvm.ir.assert_structural_equal(ret.struct_info, expected_sinfo)
+
+
+def test_conv2d_infer_struct_info():
+    bb = relax.BlockBuilder()
+    x0 = relax.Var("x", R.Tensor((2, 3, 28, 28), "float32"))
+    x1 = relax.Var("x", R.Tensor((2, 28, 28, 3), "float32"))
+    x2 = relax.Var("x", R.Tensor("float32", ndim=4))
+    x3 = relax.Var("x", R.Tensor("float32"))
+    x4 = relax.Var("x", R.Tensor())
+    x5 = relax.Var("x", R.Tensor((2, 4, 28, 28, 16), "float32"))
+    w0 = relax.Var("w", R.Tensor((4, 3, 3, 3), "float32"))
+    w1 = relax.Var("w", R.Tensor((3, 4, 3, 3), "float32"))
+    w2 = relax.Var("w", R.Tensor("float32", ndim=4))
+    w3 = relax.Var("w", R.Tensor("float32"))
+    w4 = relax.Var("w", R.Tensor((48, 4, 3, 3, 16), "float32"))
+
+    _check_inference(
+        bb, relax.op.nn.conv2d(x0, w0), relax.TensorStructInfo((2, 4, 26, 26), 
"float32")
+    )
+    _check_inference(
+        bb,
+        relax.op.nn.conv2d(x0, w0, out_dtype="float16"),
+        relax.TensorStructInfo((2, 4, 26, 26), "float16"),
+    )
+    _check_inference(
+        bb, relax.op.nn.conv2d(x0, w0, padding=1), relax.TensorStructInfo((2, 
4, 28, 28), "float32")
+    )
+    _check_inference(
+        bb,
+        relax.op.nn.conv2d(x0, w0, padding=[1, 2]),
+        relax.TensorStructInfo((2, 4, 28, 30), "float32"),
+    )
+    _check_inference(
+        bb,
+        relax.op.nn.conv2d(x0, w0, padding=[1, 2, 3, 4]),
+        relax.TensorStructInfo((2, 4, 30, 32), "float32"),
+    )
+    _check_inference(
+        bb,
+        relax.op.nn.conv2d(x0, w0, strides=2),
+        relax.TensorStructInfo((2, 4, 13, 13), "float32"),
+    )
+    _check_inference(
+        bb,
+        relax.op.nn.conv2d(x0, w0, strides=(2, 3)),
+        relax.TensorStructInfo((2, 4, 13, 9), "float32"),
+    )
+    _check_inference(
+        bb,
+        relax.op.nn.conv2d(x0, w0, dilation=2),
+        relax.TensorStructInfo((2, 4, 24, 24), "float32"),
+    )
+    _check_inference(
+        bb,
+        relax.op.nn.conv2d(x0, w0, dilation=(2, 1)),
+        relax.TensorStructInfo((2, 4, 24, 26), "float32"),
+    )
+    _check_inference(
+        bb,
+        relax.op.nn.conv2d(x1, w0, data_layout="NHWC"),
+        relax.TensorStructInfo((2, 26, 26, 4), "float32"),
+    )
+    _check_inference(
+        bb,
+        relax.op.nn.conv2d(x0, w0, out_layout="NHWC"),
+        relax.TensorStructInfo((2, 26, 26, 4), "float32"),
+    )
+    _check_inference(
+        bb,
+        relax.op.nn.conv2d(x0, w1, kernel_layout="IOHW"),
+        relax.TensorStructInfo((2, 4, 26, 26), "float32"),
+    )
+    _check_inference(
+        bb,
+        relax.op.nn.conv2d(
+            x5, w4, data_layout="NCHW16c", kernel_layout="OIHW16i", 
out_layout="NHWC16c"
+        ),
+        relax.TensorStructInfo((2, 26, 26, 3, 16), "float32"),
+    )
+    _check_inference(
+        bb, relax.op.nn.conv2d(x2, w0), 
relax.TensorStructInfo(dtype="float32", ndim=4)
+    )
+    _check_inference(
+        bb, relax.op.nn.conv2d(x3, w0), 
relax.TensorStructInfo(dtype="float32", ndim=4)
+    )
+    _check_inference(
+        bb, relax.op.nn.conv2d(x0, w2), 
relax.TensorStructInfo(dtype="float32", ndim=4)
+    )
+    _check_inference(
+        bb, relax.op.nn.conv2d(x0, w3), 
relax.TensorStructInfo(dtype="float32", ndim=4)
+    )
+    _check_inference(bb, relax.op.nn.conv2d(x4, w0), 
relax.TensorStructInfo(dtype="", ndim=4))
+
+
+def test_conv2d_infer_struct_info_shape_symbolic():
+    bb = relax.BlockBuilder()
+    n = tir.Var("n", "int64")
+    c = tir.Var("c", "int64")
+    c16 = tir.Var("c16", "int64")
+    ih = tir.Var("ih", "int64")
+    iw = tir.Var("iw", "int64")
+    ki = tir.Var("ki", "int64")
+    ko = tir.Var("ko", "int64")
+    kh = tir.Var("kh", "int64")
+    kw = tir.Var("kw", "int64")
+    x0 = relax.Var("x", R.Tensor((n, c, ih, iw), "float32"))
+    x1 = relax.Var("x", R.Tensor((n, c, ih, iw, c16), "float32"))
+    w0 = relax.Var("w", R.Tensor((ko, ki, kh, kw), "float32"))
+    w1 = relax.Var("w", R.Tensor((ko, c, kh, kw), "float32"))
+    w2 = relax.Var("w", R.Tensor((ko, c, kh, kw, c16), "float32"))
+
+    _check_inference(
+        bb,
+        relax.op.nn.conv2d(x0, w0),
+        relax.TensorStructInfo((n, ko, ih + 1 - kh, iw + 1 - kw), "float32"),
+    )
+    _check_inference(
+        bb,
+        relax.op.nn.conv2d(x0, w1),
+        relax.TensorStructInfo((n, ko, ih + 1 - kh, iw + 1 - kw), "float32"),
+    )
+    _check_inference(
+        bb,
+        relax.op.nn.conv2d(
+            x1, w2, data_layout="NCHW16c", kernel_layout="OIHW16i", 
out_layout="NCHW"
+        ),
+        relax.TensorStructInfo((n, ko, ih + 1 - kh, iw + 1 - kw), "float32"),
+    )
+    _check_inference(
+        bb,
+        relax.op.nn.conv2d(x0, w0, strides=(2, 2), padding=(1, 1), 
dilation=(2, 2)),
+        relax.TensorStructInfo(
+            (n, ko, tvm.tir.floordiv(ih + 3, 2) + 1 - kh, tvm.tir.floordiv(iw 
+ 3, 2) + 1 - kw),
+            "float32",
+        ),
+    )
+
+
+def test_conv2d_infer_struct_info_shape_var():
+    bb = relax.BlockBuilder()
+    s0 = relax.Var("s", relax.ShapeStructInfo(ndim=4))
+    s1 = relax.Var("s", relax.ShapeStructInfo(ndim=5))
+    s2 = relax.Var("s", relax.ShapeStructInfo(ndim=4))
+    s3 = relax.Var("s", relax.ShapeStructInfo())
+    x0 = relax.Var("x", relax.TensorStructInfo(s0, "float32"))
+    x1 = relax.Var("x", relax.TensorStructInfo(s1, "float32"))
+    x2 = relax.Var("x", relax.TensorStructInfo(s3, "float32"))
+    w = relax.Var("w", relax.TensorStructInfo(s2, "float32"))
+
+    _check_inference(bb, relax.op.nn.conv2d(x0, w), 
relax.TensorStructInfo(dtype="float32", ndim=4))
+    _check_inference(
+        bb,
+        relax.op.nn.conv2d(x1, w, data_layout="NCHW16c"),
+        relax.TensorStructInfo(dtype="float32", ndim=5),
+    )
+    _check_inference(
+        bb,
+        relax.op.nn.conv2d(x0, w, out_layout="NCHW16c"),
+        relax.TensorStructInfo(dtype="float32", ndim=5),
+    )
+    _check_inference(
+        bb,
+        relax.op.nn.conv2d(x2, w),
+        relax.TensorStructInfo(dtype="float32", ndim=4),
+    )
+
+
+def test_conv2d_infer_struct_info_groups():
+    bb = relax.BlockBuilder()
+    x0 = relax.Var("x", R.Tensor((2, 128, 28, 28), "float32"))
+    x1 = relax.Var("x", R.Tensor((2, 8, 28, 28, 16), "float32"))
+    w0 = relax.Var("w", R.Tensor((48, 16, 3, 3), "float32"))
+    w1 = relax.Var("w", R.Tensor((48, 2, 3, 3, 8), "float32"))
+
+    _check_inference(
+        bb, relax.op.nn.conv2d(x0, w0, groups=8), relax.TensorStructInfo((2, 
48, 26, 26), "float32")
+    )
+    _check_inference(
+        bb,
+        relax.op.nn.conv2d(x0, w1, kernel_layout="OIHW8i", groups=8),
+        relax.TensorStructInfo((2, 48, 26, 26), "float32"),
+    )
+    _check_inference(
+        bb,
+        relax.op.nn.conv2d(x1, w0, data_layout="NCHW16c", groups=8),
+        relax.TensorStructInfo((2, 3, 26, 26, 16), "float32"),
+    )
+
+
+def test_conv2d_infer_struct_info_symbolic_groups():
+    bb = relax.BlockBuilder()
+    n = tir.Var("n", "int64")
+    ic = tir.Var("c", "int64")
+    oc = tir.Var("oc", "int64")
+    x = relax.Var("x", R.Tensor((n, ic * 4, 28, 28), "float32"))
+    w0 = relax.Var("w", R.Tensor((oc * 4, ic, 3, 3), "float32"))
+    w1 = relax.Var("w", R.Tensor((oc, ic, 3, 3), "float32"))
+
+    _check_inference(
+        bb,
+        relax.op.nn.conv2d(x, w0, groups=4),
+        relax.TensorStructInfo((n, oc * 4, 26, 26), "float32"),
+    )
+    _check_inference(
+        bb, relax.op.nn.conv2d(x, w1, groups=4), relax.TensorStructInfo((n, 
oc, 26, 26), "float32")
+    )
+
+
+def test_conv2d_infer_struct_info_input_channel_group_incompatible():
+    bb = relax.BlockBuilder()
+    n = tir.Var("n", "int64")
+    ic = tir.Var("c", "int64")
+    oc = tir.Var("oc", "int64")
+    x0 = relax.Var("x", R.Tensor((2, 128, 28, 28), "float32"))
+    w0 = relax.Var("w", R.Tensor((48, 20, 3, 3), "float32"))
+    x1 = relax.Var("x", R.Tensor((n, ic * 6, 28, 28), "float32"))
+    w1 = relax.Var("w", R.Tensor((oc, ic - 1, 3, 3), "float32"))
+
+    with pytest.raises(TVMError):
+        bb.normalize(relax.op.nn.conv2d(x0, w0, groups=6))
+    with pytest.raises(TVMError):
+        bb.normalize(relax.op.nn.conv2d(x1, w1, groups=6))
+
+
+def test_conv2d_infer_struct_info_output_channel_group_incompatible():
+    bb = relax.BlockBuilder()
+    n = tir.Var("n", "int64")
+    ic = tir.Var("c", "int64")
+    oc = tir.Var("oc", "int64")
+    x0 = relax.Var("x", R.Tensor((2, 120, 28, 28), "float32"))
+    w0 = relax.Var("w", R.Tensor((128, 20, 3, 3), "float32"))
+    x1 = relax.Var("x", R.Tensor((n, ic * 6, 28, 28), "float32"))
+    w1 = relax.Var("w", R.Tensor((oc * 6 + 4, ic * 6, 3, 3), "float32"))
+
+    with pytest.raises(TVMError):
+        bb.normalize(relax.op.nn.conv2d(x0, w0, groups=6))
+    with pytest.raises(TVMError):
+        bb.normalize(relax.op.nn.conv2d(x1, w1, groups=6))
+
+
+def test_conv2d_non_positive_group():
+    x = relax.Var("x", R.Tensor((2, 128, 28, 28), "float32"))
+    w = relax.Var("w", R.Tensor((48, 16, 3, 3), "float32"))
+
+    with pytest.raises(TVMError):
+        relax.op.nn.conv2d(x, w, groups=0)
+    with pytest.raises(TVMError):
+        relax.op.nn.conv2d(x, w, groups=-2)
+
+
+def test_conv2d_infer_struct_info_more_input_dtype():
+    bb = relax.BlockBuilder()
+    x0 = relax.Var("x", R.Tensor((2, 3, 28, 28), "float16"))
+    w0 = relax.Var("w", R.Tensor((4, 3, 3, 3), "float16"))
+    x1 = relax.Var("x", R.Tensor((2, 3, 28, 28), "float64"))
+    w1 = relax.Var("w", R.Tensor((4, 3, 3, 3), "float64"))
+    x2 = relax.Var("x", R.Tensor((2, 3, 28, 28), "int8"))
+    w2 = relax.Var("w", R.Tensor((4, 3, 3, 3), "int8"))
+    x3 = relax.Var("x", R.Tensor((2, 3, 28, 28), "int32"))
+    w3 = relax.Var("w", R.Tensor((4, 3, 3, 3), "int32"))
+
+    _check_inference(
+        bb, relax.op.nn.conv2d(x0, w0), relax.TensorStructInfo((2, 4, 26, 26), 
"float16")
+    )
+    _check_inference(
+        bb, relax.op.nn.conv2d(x1, w1), relax.TensorStructInfo((2, 4, 26, 26), 
"float64")
+    )
+    _check_inference(bb, relax.op.nn.conv2d(x2, w2), 
relax.TensorStructInfo((2, 4, 26, 26), "int8"))
+    _check_inference(
+        bb, relax.op.nn.conv2d(x3, w3), relax.TensorStructInfo((2, 4, 26, 26), 
"int32")
+    )
+
+
+def test_conv2d_infer_struct_info_mixed_precision():
+    bb = relax.BlockBuilder()
+    x0 = relax.Var("x", R.Tensor((2, 3, 28, 28), "float16"))
+    w0 = relax.Var("w", R.Tensor((4, 3, 3, 3), "float16"))
+    x1 = relax.Var("x", R.Tensor((2, 3, 28, 28), "int8"))
+    w1 = relax.Var("w", R.Tensor((4, 3, 3, 3), "int8"))
+    x2 = relax.Var("x", R.Tensor((2, 3, 28, 28)))
+    w2 = relax.Var("w", R.Tensor((4, 3, 3, 3)))
+
+    _check_inference(
+        bb,
+        relax.op.nn.conv2d(x0, w0, out_dtype="float32"),
+        relax.TensorStructInfo((2, 4, 26, 26), "float32"),
+    )
+    _check_inference(
+        bb,
+        relax.op.nn.conv2d(x1, w1, out_dtype="int32"),
+        relax.TensorStructInfo((2, 4, 26, 26), "int32"),
+    )
+    _check_inference(
+        bb,
+        relax.op.nn.conv2d(x2, w2, out_dtype="float32"),
+        relax.TensorStructInfo((2, 4, 26, 26), "float32"),
+    )
+
+
+def test_conv2d_unequal_input_channel():
+    bb = relax.BlockBuilder()
+    ic = tir.Var("ic", "int64")
+    x0 = relax.Var("x", R.Tensor([2, 3, 28, 28], "float32"))
+    w0 = relax.Var("w", R.Tensor([3, 4, 3, 3], "float32"))
+    x1 = relax.Var("x", R.Tensor([2, ic, 28, 28], "float32"))
+    w1 = relax.Var("w", R.Tensor([4, ic + 2, 3, 3], "float32"))
+    with pytest.raises(TVMError):
+        bb.normalize(relax.op.nn.conv2d(x0, w0))
+    with pytest.raises(TVMError):
+        bb.normalize(relax.op.nn.conv2d(x1, w1))
+
+
+def test_conv2d_stride_padding_dilation_int64():
+    x = relax.Var("x", R.Tensor((2, 3, 28, 28), "float32"))
+    w = relax.Var("w", R.Tensor((4, 3, 3, 3), "float32"))
+    conv2d = relax.op.nn.conv2d(x, w, strides=(1, 1), padding=(1, 1), 
dilation=(1, 1))
+
+    assert conv2d.attrs.strides[0].dtype == "int64"
+    assert conv2d.attrs.strides[1].dtype == "int64"
+    assert conv2d.attrs.padding[0].dtype == "int64"
+    assert conv2d.attrs.padding[1].dtype == "int64"
+    assert conv2d.attrs.padding[2].dtype == "int64"
+    assert conv2d.attrs.padding[3].dtype == "int64"
+    assert conv2d.attrs.dilation[0].dtype == "int64"
+    assert conv2d.attrs.dilation[1].dtype == "int64"
+
+
+def test_conv2d_wrong_strides_padding_dilation_length():
+    x = relax.Var("x", R.Tensor((2, 3, 28, 28), "float32"))
+    w = relax.Var("w", R.Tensor((4, 3, 3, 3), "float32"))
+    with pytest.raises(TVMError):
+        relax.op.nn.conv2d(x, w, strides=(1, 2, 3))
+    with pytest.raises(TVMError):
+        relax.op.nn.conv2d(x, w, padding=(1, 2, 3))
+    with pytest.raises(TVMError):
+        relax.op.nn.conv2d(x, w, dilation=(1, 2, 3))
+
+
+def test_conv2d_infer_struct_info_wrong_layout_string():
+    bb = relax.BlockBuilder()
+    x = relax.Var("x", R.Tensor((2, 3, 28, 28), "float32"))
+    w = relax.Var("w", R.Tensor((4, 3, 3, 3), "float32"))
+    with pytest.raises(TVMError):
+        bb.normalize(relax.op.nn.conv2d(x, w, data_layout="OIHW"))
+    with pytest.raises(TVMError):
+        bb.normalize(relax.op.nn.conv2d(x, w, kernel_layout="NHWC"))
+    with pytest.raises(TVMError):
+        bb.normalize(relax.op.nn.conv2d(x, w, out_layout="OHWI"))
+
+
+def test_conv2d_dtype_mismatch():
+    bb = relax.BlockBuilder()
+    x = relax.Var("x", R.Tensor((2, 3, 28, 28), "float32"))
+    w = relax.Var("w", R.Tensor((4, 3, 3, 3), "int8"))
+
+    with pytest.raises(TVMError):
+        bb.normalize(relax.op.nn.conv2d(x, w))
+
+
+def test_conv2d_wrong_input_ndim():
+    bb = relax.BlockBuilder()
+    x0 = relax.Var("x", R.Tensor((2, 3, 28, 28), "float32"))
+    x1 = relax.Var("x", R.Tensor((2, 3, 28, 28, 3), "float32"))
+    x2 = relax.Var("x", R.Tensor("float32", ndim=3))
+    w0 = relax.Var("w", R.Tensor((4, 3, 3, 3), "float32"))
+    w1 = relax.Var("w", R.Tensor((4, 3, 6, 3, 3), "float32"))
+    w2 = relax.Var("w", R.Tensor("float32", ndim=6))
+
+    with pytest.raises(TVMError):
+        bb.normalize(relax.op.nn.conv2d(x0, w1))
+    with pytest.raises(TVMError):
+        bb.normalize(relax.op.nn.conv2d(x0, w1, data_layout="NCHW16c"))
+    with pytest.raises(TVMError):
+        bb.normalize(relax.op.nn.conv2d(x0, w2))
+    with pytest.raises(TVMError):
+        bb.normalize(relax.op.nn.conv2d(x1, w0))
+    with pytest.raises(TVMError):
+        bb.normalize(relax.op.nn.conv2d(x2, w0))
+
+
+def test_conv2d_infer_struct_info_wrong_input_type():
+    bb = relax.BlockBuilder()
+    x0 = relax.Var("x", R.Tensor((2, 3, 28, 28), "float32"))
+    x1 = relax.Var("x", relax.ShapeStructInfo((2, 3, 28, 28)))
+    w0 = relax.Var("w", R.Tensor((4, 3, 3, 3), "float32"))
+    w1 = relax.Var("w", relax.FuncStructInfo([], R.Tensor((4, 3, 3, 3), 
"float32")))
+
+    with pytest.raises(TVMError):
+        bb.normalize(relax.op.nn.conv2d(x0, w1))
+    with pytest.raises(TVMError):
+        bb.normalize(relax.op.nn.conv2d(x1, w0))
+
+
+if __name__ == "__main__":
+    tvm.testing.main()
diff --git a/tests/python/relax/test_op_nn_pooling.py 
b/tests/python/relax/test_op_nn_pooling.py
new file mode 100644
index 0000000000..0eec5de21c
--- /dev/null
+++ b/tests/python/relax/test_op_nn_pooling.py
@@ -0,0 +1,429 @@
+# 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.
+import pytest
+import tvm
+import tvm.testing
+from tvm import relax, tir
+from tvm import TVMError
+from tvm.ir import Op
+from tvm.script import relax as R
+
+
+def test_op_correctness():
+    x = relax.Var("x", R.Tensor((2, 3, 28, 28), "float32"))
+    assert relax.op.nn.max_pool2d(x).op == Op.get("relax.nn.max_pool2d")
+    assert relax.op.nn.adaptive_avg_pool2d(x).op == 
Op.get("relax.nn.adaptive_avg_pool2d")
+
+
+def _check_inference(bb: relax.BlockBuilder, call: relax.Call, expected_sinfo: 
relax.StructInfo):
+    ret = bb.normalize(call)
+    tvm.ir.assert_structural_equal(ret.struct_info, expected_sinfo)
+
+
+def test_max_pool2d_infer_struct_info():
+    bb = relax.BlockBuilder()
+    x0 = relax.Var("x", R.Tensor((2, 3, 32, 32), "float32"))
+    x1 = relax.Var("x", R.Tensor((2, 32, 32, 3), "float32"))
+    x2 = relax.Var("x", R.Tensor("float32", ndim=4))
+    x3 = relax.Var("x", R.Tensor("float32"))
+    x4 = relax.Var("x", R.Tensor(ndim=4))
+    x5 = relax.Var("x", R.Tensor())
+    x6 = relax.Var("x", R.Tensor((2, 4, 32, 32, 16), "float32"))
+
+    _check_inference(
+        bb, relax.op.nn.max_pool2d(x0), relax.TensorStructInfo((2, 3, 32, 32), 
"float32")
+    )
+    _check_inference(
+        bb,
+        relax.op.nn.max_pool2d(x0, pool_size=3),
+        relax.TensorStructInfo((2, 3, 30, 30), "float32"),
+    )
+    _check_inference(
+        bb,
+        relax.op.nn.max_pool2d(x0, pool_size=(5, 3)),
+        relax.TensorStructInfo((2, 3, 28, 30), "float32"),
+    )
+    _check_inference(
+        bb, relax.op.nn.max_pool2d(x0, padding=1), relax.TensorStructInfo((2, 
3, 34, 34), "float32")
+    )
+    _check_inference(
+        bb,
+        relax.op.nn.max_pool2d(x0, padding=[1, 2]),
+        relax.TensorStructInfo((2, 3, 34, 36), "float32"),
+    )
+    _check_inference(
+        bb,
+        relax.op.nn.max_pool2d(x0, strides=2),
+        relax.TensorStructInfo((2, 3, 16, 16), "float32"),
+    )
+    _check_inference(
+        bb,
+        relax.op.nn.max_pool2d(x0, dilation=2),
+        relax.TensorStructInfo((2, 3, 32, 32), "float32"),
+    )
+    _check_inference(
+        bb,
+        relax.op.nn.max_pool2d(x1, layout="NHWC"),
+        relax.TensorStructInfo((2, 32, 32, 3), "float32"),
+    )
+    _check_inference(
+        bb,
+        relax.op.nn.max_pool2d(x0, out_layout="NHWC"),
+        relax.TensorStructInfo((2, 32, 32, 3), "float32"),
+    )
+    _check_inference(
+        bb,
+        relax.op.nn.max_pool2d(x6, layout="NCHW16c", out_layout="NHWC16c"),
+        relax.TensorStructInfo((2, 32, 32, 4, 16), "float32"),
+    )
+    _check_inference(
+        bb, relax.op.nn.max_pool2d(x2), 
relax.TensorStructInfo(dtype="float32", ndim=4)
+    )
+    _check_inference(
+        bb, relax.op.nn.max_pool2d(x3), 
relax.TensorStructInfo(dtype="float32", ndim=4)
+    )
+    _check_inference(bb, relax.op.nn.max_pool2d(x4), 
relax.TensorStructInfo(dtype="", ndim=4))
+    _check_inference(bb, relax.op.nn.max_pool2d(x5), 
relax.TensorStructInfo(dtype="", ndim=4))
+
+
+def test_max_pool2d_infer_struct_info_shape_symbolic():
+    bb = relax.BlockBuilder()
+    n = tir.Var("n", "int64")
+    c = tir.Var("c", "int64")
+    c16 = tir.Var("c16", "int64")
+    ih = tir.Var("ih", "int64")
+    iw = tir.Var("iw", "int64")
+    x0 = relax.Var("x", R.Tensor((n, c, ih, iw), "float32"))
+    x1 = relax.Var("x", R.Tensor((n, c, ih, iw, c16), "float32"))
+
+    _check_inference(
+        bb,
+        relax.op.nn.max_pool2d(
+            x0, pool_size=(3, 3), strides=(3, 3), padding=(2, 2), dilation=(2, 
2)
+        ),
+        relax.TensorStructInfo(
+            (
+                n,
+                c,
+                tvm.tir.floordiv(ih - 1, 3) + 1,
+                tvm.tir.floordiv(iw - 1, 3) + 1,
+            ),
+            "float32",
+        ),
+    )
+    _check_inference(
+        bb,
+        relax.op.nn.max_pool2d(x1, layout="NCHW16c", out_layout="NHWC"),
+        relax.TensorStructInfo((n, ih, iw, c * 16), "float32"),
+    )
+
+
+def test_max_pool2d_infer_struct_info_shape_var():
+    bb = relax.BlockBuilder()
+    s0 = relax.Var("s", relax.ShapeStructInfo(ndim=4))
+    s1 = relax.Var("s", relax.ShapeStructInfo(ndim=5))
+    s2 = relax.Var("s", relax.ShapeStructInfo())
+    x0 = relax.Var("x", relax.TensorStructInfo(s0, "float32"))
+    x1 = relax.Var("x", relax.TensorStructInfo(s1, "float32"))
+    x2 = relax.Var("x", relax.TensorStructInfo(s2, "float32"))
+
+    _check_inference(
+        bb, relax.op.nn.max_pool2d(x0), 
relax.TensorStructInfo(dtype="float32", ndim=4)
+    )
+    _check_inference(
+        bb,
+        relax.op.nn.max_pool2d(x1, layout="NCHW16c"),
+        relax.TensorStructInfo(dtype="float32", ndim=5),
+    )
+    _check_inference(
+        bb,
+        relax.op.nn.max_pool2d(x2),
+        relax.TensorStructInfo(dtype="float32", ndim=4),
+    )
+
+
+def test_max_pool2d_infer_struct_info_ceil_mode():
+    bb = relax.BlockBuilder()
+    x = relax.Var("x", R.Tensor((2, 3, 32, 32), "float32"))
+
+    _check_inference(
+        bb,
+        relax.op.nn.max_pool2d(x, pool_size=3, strides=2, ceil_mode=True),
+        relax.TensorStructInfo((2, 3, 16, 16), "float32"),
+    )
+    _check_inference(
+        bb,
+        relax.op.nn.max_pool2d(x, pool_size=(5, 3), strides=2, ceil_mode=True),
+        relax.TensorStructInfo((2, 3, 15, 16), "float32"),
+    )
+
+
+def test_max_pool2d_infer_struct_info_ceil_mode_symbolic():
+    bb = relax.BlockBuilder()
+    n = tir.Var("n", "int64")
+    c = tir.Var("c", "int64")
+    ih = tir.Var("ih", "int64")
+    iw = tir.Var("iw", "int64")
+    x = relax.Var("x", R.Tensor((n, c, ih, iw), "float32"))
+
+    _check_inference(
+        bb,
+        relax.op.nn.max_pool2d(
+            x, pool_size=(3, 3), strides=(2, 2), padding=(1, 1), dilation=(2, 
2), ceil_mode=True
+        ),
+        relax.TensorStructInfo((n, c, tvm.tir.floordiv(ih, 2), 
tvm.tir.floordiv(iw, 2)), "float32"),
+    )
+
+
+def test_max_pool2d_infer_struct_info_more_input_dtype():
+    bb = relax.BlockBuilder()
+    x0 = relax.Var("x", R.Tensor((2, 3, 32, 32), "float16"))
+    x1 = relax.Var("x", R.Tensor((2, 3, 32, 32), "int8"))
+    x2 = relax.Var("x", R.Tensor((2, 3, 32, 32), "int64"))
+    _check_inference(
+        bb, relax.op.nn.max_pool2d(x0), relax.TensorStructInfo((2, 3, 32, 32), 
"float16")
+    )
+    _check_inference(bb, relax.op.nn.max_pool2d(x1), 
relax.TensorStructInfo((2, 3, 32, 32), "int8"))
+    _check_inference(
+        bb, relax.op.nn.max_pool2d(x2), relax.TensorStructInfo((2, 3, 32, 32), 
"int64")
+    )
+
+
+def test_conv2d_stride_padding_dilation_int64():
+    x = relax.Var("x", R.Tensor((2, 3, 28, 28), "float32"))
+    max_pool2d = relax.op.nn.max_pool2d(x, (3, 3), strides=(1, 1), padding=(1, 
1), dilation=(1, 1))
+
+    assert max_pool2d.attrs.strides[0].dtype == "int64"
+    assert max_pool2d.attrs.strides[1].dtype == "int64"
+    assert max_pool2d.attrs.padding[0].dtype == "int64"
+    assert max_pool2d.attrs.padding[1].dtype == "int64"
+    assert max_pool2d.attrs.padding[2].dtype == "int64"
+    assert max_pool2d.attrs.padding[3].dtype == "int64"
+    assert max_pool2d.attrs.dilation[0].dtype == "int64"
+    assert max_pool2d.attrs.dilation[1].dtype == "int64"
+
+
+def test_max_pool2d_wrong_pool_size_strides_padding_dilation_length():
+    x = relax.Var("x", R.Tensor((2, 3, 28, 28), "float32"))
+    with pytest.raises(TVMError):
+        relax.op.nn.max_pool2d(x, pool_size=(1, 2, 3))
+    with pytest.raises(TVMError):
+        relax.op.nn.max_pool2d(x, strides=(1, 2, 3))
+    with pytest.raises(TVMError):
+        relax.op.nn.max_pool2d(x, padding=(1, 2, 3))
+    with pytest.raises(TVMError):
+        relax.op.nn.max_pool2d(x, dilation=(1, 2, 3))
+
+
+def test_max_pool2d_infer_struct_info_wrong_layout_string():
+    bb = relax.BlockBuilder()
+    x = relax.Var("x", R.Tensor((2, 3, 28, 28), "float32"))
+    with pytest.raises(TVMError):
+        bb.normalize(relax.op.nn.max_pool2d(x, layout="OIHW"))
+    with pytest.raises(TVMError):
+        bb.normalize(relax.op.nn.max_pool2d(x, out_layout="OHWI"))
+
+
+def test_max_pool2d_wrong_input_ndim():
+    bb = relax.BlockBuilder()
+    x0 = relax.Var("x", R.Tensor((2, 3, 28, 28, 3), "float32"))
+    x1 = relax.Var("x", R.Tensor("float32", ndim=3))
+    with pytest.raises(TVMError):
+        bb.normalize(relax.op.nn.max_pool2d(x0))
+    with pytest.raises(TVMError):
+        bb.normalize(relax.op.nn.max_pool2d(x1))
+
+
+def test_max_pool2d_infer_struct_info_wrong_input_type():
+    bb = relax.BlockBuilder()
+    x0 = relax.Var("x", relax.ShapeStructInfo((2, 3, 28, 28)))
+    x1 = relax.Var("x", relax.FuncStructInfo([], R.Tensor((2, 3, 28, 28), 
"float32")))
+
+    with pytest.raises(TVMError):
+        bb.normalize(relax.op.nn.max_pool2d(x0))
+    with pytest.raises(TVMError):
+        bb.normalize(relax.op.nn.max_pool2d(x1))
+
+
+def test_adaptive_avg_pool2d_infer_struct_info():
+    bb = relax.BlockBuilder()
+    x0 = relax.Var("x", R.Tensor((2, 3, 32, 32), "float32"))
+    x1 = relax.Var("x", R.Tensor((2, 32, 32, 3), "float32"))
+    x2 = relax.Var("x", R.Tensor("float32", ndim=4))
+    x3 = relax.Var("x", R.Tensor("float32"))
+    x4 = relax.Var("x", R.Tensor(ndim=4))
+    x5 = relax.Var("x", R.Tensor())
+    x6 = relax.Var("x", R.Tensor((2, 4, 32, 32, 16), "float32"))
+
+    _check_inference(
+        bb, relax.op.nn.adaptive_avg_pool2d(x0), relax.TensorStructInfo((2, 3, 
32, 32), "float32")
+    )
+    _check_inference(
+        bb,
+        relax.op.nn.adaptive_avg_pool2d(x0, output_size=30),
+        relax.TensorStructInfo((2, 3, 30, 30), "float32"),
+    )
+    _check_inference(
+        bb,
+        relax.op.nn.adaptive_avg_pool2d(x0, output_size=(28, 30)),
+        relax.TensorStructInfo((2, 3, 28, 30), "float32"),
+    )
+    _check_inference(
+        bb,
+        relax.op.nn.adaptive_avg_pool2d(x1, layout="NHWC"),
+        relax.TensorStructInfo((2, 32, 32, 3), "float32"),
+    )
+    _check_inference(
+        bb,
+        relax.op.nn.adaptive_avg_pool2d(x0, out_layout="NHWC"),
+        relax.TensorStructInfo((2, 32, 32, 3), "float32"),
+    )
+    _check_inference(
+        bb,
+        relax.op.nn.adaptive_avg_pool2d(x6, layout="NCHW16c", 
out_layout="NHWC16c"),
+        relax.TensorStructInfo((2, 32, 32, 4, 16), "float32"),
+    )
+    _check_inference(
+        bb, relax.op.nn.adaptive_avg_pool2d(x2), 
relax.TensorStructInfo(dtype="float32", ndim=4)
+    )
+    _check_inference(
+        bb, relax.op.nn.adaptive_avg_pool2d(x3), 
relax.TensorStructInfo(dtype="float32", ndim=4)
+    )
+    _check_inference(
+        bb, relax.op.nn.adaptive_avg_pool2d(x4), 
relax.TensorStructInfo(dtype="", ndim=4)
+    )
+    _check_inference(
+        bb, relax.op.nn.adaptive_avg_pool2d(x5), 
relax.TensorStructInfo(dtype="", ndim=4)
+    )
+
+
+def test_adaptive_avg_pool2d_infer_struct_info_shape_symbolic():
+    bb = relax.BlockBuilder()
+    n = tir.Var("n", "int64")
+    c = tir.Var("c", "int64")
+    c16 = tir.Var("c16", "int64")
+    ih = tir.Var("ih", "int64")
+    iw = tir.Var("iw", "int64")
+    x0 = relax.Var("x", R.Tensor((n, c, ih, iw), "float32"))
+    x1 = relax.Var("x", R.Tensor((n, c, ih, iw, c16), "float32"))
+
+    _check_inference(
+        bb, relax.op.nn.adaptive_avg_pool2d(x0), relax.TensorStructInfo((n, c, 
ih, iw), "float32")
+    )
+    _check_inference(
+        bb,
+        relax.op.nn.adaptive_avg_pool2d(x0, output_size=256),
+        relax.TensorStructInfo((n, c, 256, 256), "float32"),
+    )
+    _check_inference(
+        bb,
+        relax.op.nn.adaptive_avg_pool2d(x0, output_size=(256, 128)),
+        relax.TensorStructInfo((n, c, 256, 128), "float32"),
+    )
+    _check_inference(
+        bb,
+        relax.op.nn.adaptive_avg_pool2d(x1, layout="NCHW16c", 
out_layout="NHWC"),
+        relax.TensorStructInfo((n, ih, iw, c * 16), "float32"),
+    )
+
+
+def test_adaptive_avg_pool2d_infer_struct_info_shape_var():
+    bb = relax.BlockBuilder()
+    s0 = relax.Var("s", relax.ShapeStructInfo(ndim=4))
+    s1 = relax.Var("s", relax.ShapeStructInfo(ndim=5))
+    s2 = relax.Var("s", relax.ShapeStructInfo())
+    x0 = relax.Var("x", relax.TensorStructInfo(s0, "float32"))
+    x1 = relax.Var("x", relax.TensorStructInfo(s1, "float32"))
+    x2 = relax.Var("x", relax.TensorStructInfo(s2, "float32"))
+
+    _check_inference(bb, relax.op.nn.adaptive_avg_pool2d(x0), 
relax.TensorStructInfo(s0, "float32"))
+    _check_inference(
+        bb,
+        relax.op.nn.adaptive_avg_pool2d(x0, output_size=32),
+        relax.TensorStructInfo(dtype="float32", ndim=4),
+    )
+    _check_inference(
+        bb,
+        relax.op.nn.adaptive_avg_pool2d(x1, layout="NCHW16c"),
+        relax.TensorStructInfo(s1, "float32"),
+    )
+    _check_inference(
+        bb,
+        relax.op.nn.adaptive_avg_pool2d(x0, out_layout="NCHW16c"),
+        relax.TensorStructInfo(dtype="float32", ndim=5),
+    )
+    _check_inference(
+        bb,
+        relax.op.nn.adaptive_avg_pool2d(x2, out_layout="NCHW16c"),
+        relax.TensorStructInfo(dtype="float32", ndim=5),
+    )
+
+
+def test_adaptive_avg_pool2d_infer_struct_info_more_input_dtype():
+    bb = relax.BlockBuilder()
+    x0 = relax.Var("x", R.Tensor((2, 3, 32, 32), "float16"))
+    x1 = relax.Var("x", R.Tensor((2, 3, 32, 32), "int8"))
+    x2 = relax.Var("x", R.Tensor((2, 3, 32, 32), "int64"))
+    _check_inference(
+        bb, relax.op.nn.adaptive_avg_pool2d(x0), relax.TensorStructInfo((2, 3, 
32, 32), "float16")
+    )
+    _check_inference(
+        bb, relax.op.nn.adaptive_avg_pool2d(x1), relax.TensorStructInfo((2, 3, 
32, 32), "int8")
+    )
+    _check_inference(
+        bb, relax.op.nn.adaptive_avg_pool2d(x2), relax.TensorStructInfo((2, 3, 
32, 32), "int64")
+    )
+
+
+def test_adaptive_avg_pool2d_wrong_output_size_ndim():
+    x = relax.Var("x", R.Tensor((2, 3, 32, 32), "float32"))
+    with pytest.raises(TVMError):
+        relax.op.nn.adaptive_avg_pool2d(x, (32, 32, 32))
+
+
+def test_adaptive_avg_pool2d_infer_struct_info_wrong_layout_string():
+    bb = relax.BlockBuilder()
+    x = relax.Var("x", R.Tensor((2, 3, 28, 28), "float32"))
+    with pytest.raises(TVMError):
+        bb.normalize(relax.op.nn.adaptive_avg_pool2d(x, layout="OIHW"))
+    with pytest.raises(TVMError):
+        bb.normalize(relax.op.nn.adaptive_avg_pool2d(x, out_layout="OHWI"))
+
+
+def test_adaptive_avg_pool2d_wrong_input_ndim():
+    bb = relax.BlockBuilder()
+    x0 = relax.Var("x", R.Tensor((2, 3, 28, 28, 3), "float32"))
+    x1 = relax.Var("x", R.Tensor("float32", ndim=3))
+    with pytest.raises(TVMError):
+        bb.normalize(relax.op.nn.adaptive_avg_pool2d(x0))
+    with pytest.raises(TVMError):
+        bb.normalize(relax.op.nn.adaptive_avg_pool2d(x1))
+
+
+def test_adaptive_avg_pool2d_infer_struct_info_wrong_input_type():
+    bb = relax.BlockBuilder()
+    x0 = relax.Var("x", relax.ShapeStructInfo((2, 3, 28, 28)))
+    x1 = relax.Var("x", relax.FuncStructInfo([], R.Tensor((2, 3, 28, 28), 
"float32")))
+
+    with pytest.raises(TVMError):
+        bb.normalize(relax.op.nn.adaptive_avg_pool2d(x0))
+    with pytest.raises(TVMError):
+        bb.normalize(relax.op.nn.adaptive_avg_pool2d(x1))
+
+
+if __name__ == "__main__":
+    tvm.testing.main()
diff --git a/tests/python/relax/test_tvmscript_parser_op_nn.py 
b/tests/python/relax/test_tvmscript_parser_op_nn.py
new file mode 100644
index 0000000000..4e52bccb86
--- /dev/null
+++ b/tests/python/relax/test_tvmscript_parser_op_nn.py
@@ -0,0 +1,193 @@
+# 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.
+
+from typing import Optional, Union
+
+import tvm
+import tvm.script
+import tvm.testing
+from tvm import IRModule, relax
+from tvm.script import relax as R
+
+
+def _check(
+    parsed: Union[relax.Function, IRModule],
+    expect: Optional[Union[relax.Function, IRModule]],
+):
+    test = parsed.script(show_meta=True)
+    roundtrip_mod = tvm.script.from_source(test)
+    tvm.ir.assert_structural_equal(parsed, roundtrip_mod)
+    if expect:
+        tvm.ir.assert_structural_equal(parsed, expect)
+
+
+def test_conv2d():
+    @R.function
+    def foo(
+        x: R.Tensor((2, 3, 228, 228), "float32"), w: R.Tensor((16, 3, 5, 5), 
"float32")
+    ) -> R.Tensor((2, 16, 224, 224), "float16"):
+        gv: R.Tensor((2, 16, 224, 224), "float16") = R.nn.conv2d(x, w, 
out_dtype="float16")
+        return gv
+
+    x = relax.Var("x", R.Tensor([2, 3, 228, 228], "float32"))
+    w = relax.Var("w", R.Tensor([16, 3, 5, 5], "float32"))
+    bb = relax.BlockBuilder()
+    with bb.function("foo", [x, w]):
+        gv = bb.emit(relax.op.nn.conv2d(x, w, out_dtype="float16"))
+        bb.emit_func_output(gv)
+
+    _check(foo, bb.get()["foo"])
+
+
+def test_max_pool2d():
+    @R.function
+    def foo(
+        x: R.Tensor((1, 1, 32, 32), dtype="float32")
+    ) -> R.Tensor((1, 1, 30, 30), dtype="float32"):
+        gv: R.Tensor((1, 1, 30, 30), dtype="float32") = R.nn.max_pool2d(x, 
pool_size=(3,))
+        return gv
+
+    x = relax.Var("x", R.Tensor([1, 1, 32, 32], "float32"))
+    bb = relax.BlockBuilder()
+    with bb.function("foo", [x]):
+        gv = bb.emit(relax.op.nn.max_pool2d(x, pool_size=(3,)))
+        bb.emit_func_output(gv)
+
+    _check(foo, bb.get()["foo"])
+
+
+def test_adaptive_avg_pool2d():
+    @R.function
+    def foo(x: R.Tensor((2, 64, 8, 9), "float32")) -> R.Tensor((2, 64, 7, 7), 
"float32"):
+        gv: R.Tensor((2, 64, 7, 7), "float32") = R.nn.adaptive_avg_pool2d(x, 
output_size=(7, 7))
+        return gv
+
+    x = relax.Var("x", R.Tensor((2, 64, 8, 9), dtype="float32"))
+    bb = relax.BlockBuilder()
+    with bb.function("foo", [x]):
+        gv = bb.emit(relax.op.nn.adaptive_avg_pool2d(x, output_size=(7, 7)))
+        bb.emit_func_output(gv)
+
+    _check(foo, bb.get()["foo"])
+
+
+def test_gelu():
+    @R.function
+    def foo(x: R.Tensor((2, 3), "float32")) -> R.Tensor((2, 3), "float32"):
+        gv: R.Tensor((2, 3), "float32") = R.nn.gelu(x)
+        return gv
+
+    x = relax.Var("x", R.Tensor((2, 3), "float32"))
+    bb = relax.BlockBuilder()
+    with bb.function("foo", [x]):
+        gv = bb.emit(relax.op.nn.gelu(x))
+        bb.emit_func_output(gv)
+
+    _check(foo, bb.get()["foo"])
+
+
+def test_softmax():
+    @R.function
+    def foo(x: R.Tensor((2, 3), "float32")) -> R.Tensor((2, 3), "float32"):
+        gv: R.Tensor((2, 3), "float32") = R.nn.softmax(x)
+        return gv
+
+    x = relax.Var("x", R.Tensor((2, 3), "float32"))
+    bb = relax.BlockBuilder()
+    with bb.function("foo", [x]):
+        gv = bb.emit(relax.op.nn.softmax(x))
+        bb.emit_func_output(gv)
+
+    _check(foo, bb.get()["foo"])
+
+
+def test_batch_norm():
+    @R.function
+    def foo(
+        x: R.Tensor((2, 4, 3, 3), dtype="float32"),
+        gamma: R.Tensor((4,), dtype="float32"),
+        beta: R.Tensor((4,), dtype="float32"),
+        moving_mean: R.Tensor((4,), dtype="float32"),
+        moving_var: R.Tensor((4,), dtype="float32"),
+    ) -> R.Tuple(
+        R.Tensor((2, 4, 3, 3), dtype="float32"),
+        R.Tensor((4,), dtype="float32"),
+        R.Tensor((4,), dtype="float32"),
+    ):
+        gv: R.Tuple(
+            R.Tensor((2, 4, 3, 3), dtype="float32"),
+            R.Tensor((4,), dtype="float32"),
+            R.Tensor((4,), dtype="float32"),
+        ) = R.nn.batch_norm(x, gamma, beta, moving_mean, moving_var, axis=1)
+        return gv
+
+    x = relax.Var("x", R.Tensor((2, 4, 3, 3), "float32"))
+    gamma = relax.Var("gamma", R.Tensor((4,), "float32"))
+    beta = relax.Var("beta", R.Tensor((4,), "float32"))
+    moving_mean = relax.Var("moving_mean", R.Tensor((4,), "float32"))
+    moving_var = relax.Var("moving_var", R.Tensor((4,), "float32"))
+
+    bb = relax.BlockBuilder()
+    with bb.function("foo", [x, gamma, beta, moving_mean, moving_var]):
+        gv = bb.emit(relax.op.nn.batch_norm(x, gamma, beta, moving_mean, 
moving_var, axis=1))
+        bb.emit_func_output(gv)
+
+    _check(foo, bb.get()["foo"])
+
+
+def test_layer_norm():
+    @R.function
+    def foo(
+        x: R.Tensor((2, 3, 4, 5), "float32"),
+        gamma: R.Tensor((4, 5), "float32"),
+        beta: R.Tensor((4, 5), "float32"),
+    ) -> R.Tensor((2, 3, 4, 5), "float32"):
+        gv: R.Tensor((2, 3, 4, 5), "float32") = R.nn.layer_norm(x, gamma, 
beta, axes=[-2, -1])
+        return gv
+
+    x = relax.Var("x", R.Tensor((2, 3, 4, 5), "float32"))
+    gamma = relax.Var("gamma", R.Tensor((4, 5), "float32"))
+    beta = relax.Var("beta", R.Tensor((4, 5), "float32"))
+    bb = relax.BlockBuilder()
+    with bb.function("foo", [x, gamma, beta]):
+        gv = bb.emit(relax.op.nn.layer_norm(x, gamma, beta, axes=[-2, -1]))
+        bb.emit_func_output(gv)
+
+    _check(foo, bb.get()["foo"])
+
+
+def test_dropout():
+    @R.function
+    def foo(
+        x: R.Tensor((2, 3), "float32")
+    ) -> R.Tuple(R.Tensor((2, 3), "float32"), R.Tensor((2, 3), "float32")):
+        gv: R.Tuple(R.Tensor((2, 3), "float32"), R.Tensor((2, 3), "float32")) 
= R.nn.dropout(
+            x, rate=0.5
+        )
+        return gv
+
+    x = relax.Var("x", R.Tensor((2, 3), "float32"))
+    bb = relax.BlockBuilder()
+    with bb.function("foo", [x]):
+        gv = bb.emit(relax.op.nn.dropout(x))
+        bb.emit_func_output(gv)
+
+    _check(foo, bb.get()["foo"])
+
+
+if __name__ == "__main__":
+    tvm.testing.main()