text: # 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.
  # pylint: disable=invalid-name, too-many-lines
  """Neural network operations."""
  from tvm.relay import expr
  
  from . import _make
  from ..dyn.nn import _make as _dyn_make
  from .utils import get_pad_tuple1d, get_pad_tuple2d, get_pad_tuple3d
  from ...expr import const, Expr
  
  
  def conv1d(
      data,
      weight,
      strides=1,
      padding=0,
      dilation=1,
      groups=1,
      channels=None,
      kernel_size=None,
      data_layout="NCW",
      kernel_layout="OIW",
      out_layout="",
      out_dtype="",
  ):
      r"""1D 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 `NCW`
      and kernel_layout is `OIW`, conv1d takes in
      a data Tensor with shape `(batch_size, in_channels, width)`,
      and a weight Tensor with shape `(channels, in_channels, kernel_size)`
      to produce an output Tensor with the following rule:
  
      .. math::
  
          \mbox{out}[b, c, w] = \sum_{dw, k}
             \mbox{data}[b, k, \mbox{strides}[0] * w + dw] *
             \mbox{weight}[c, k, dw]
  
      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
      (`NCW` for data and `OIW` for weight), perform the computation,
      then convert to the out_layout.
  
  
      Parameters
      ----------
      data : tvm.relay.Expr
          The input data to the operator.
  
      weight : tvm.relay.Expr
          The weight expressions.
  
      strides : Optional[int, Tuple[int]]
          The strides of convolution.
  
      padding : Optional[int, Tuple[int]]
          The padding of convolution on both sides of the input before convolution.
  
      dilation : Optional[int, Tuple[int]]
          Specifies the dilation rate to be used for dilated convolution.
  
      groups : Optional[int]
          Currently unused for 1D convolution.
  
      channels : Optional[int]
          Number of output channels of this convolution.
  
      kernel_size : Optional[int, Tuple[int]]
          The spatial dimension of the convolution kernel.
  
      data_layout : Optional[str]
          Layout of the input.
  
      kernel_layout : Optional[str]
          Layout of the weight.
  
      out_layout : Optional[str]
          Layout of the output, by default, out_layout is the same as data_layout
  
      out_dtype : Optional[str]
          Specifies the output data type for mixed precision conv2d.
  
      Returns
      -------
      result : tvm.relay.Expr
          The computed result.
      """
      if isinstance(kernel_size, int):
          kernel_size = (kernel_size,)
      if isinstance(strides, int):
          strides = (strides,)
      if isinstance(dilation, int):
          dilation = (dilation,)
      padding = get_pad_tuple1d(padding)
      return _make.conv1d(
          data,
          weight,
          strides,
          padding,
          dilation,
          groups,
          channels,
          kernel_size,
          data_layout,
          kernel_layout,
          out_layout,
          out_dtype,
      )
  
  
  def conv2d(
      data,
      weight,
      strides=(1, 1),
      padding=(0, 0),
      dilation=(1, 1),
      groups=1,
      channels=None,
      kernel_size=None,
      data_layout="NCHW",
      kernel_layout="OIHW",
      out_layout="",
      out_dtype="",
  ):
      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_size[0], kernel_size[1])`
      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 : tvm.relay.Expr
          The input data to the operator.
  
      weight : tvm.relay.Expr
          The weight expressions.
  
      strides : Optional[int, Tuple[int]]
          The strides of convolution.
  
      padding : Optional[int, Tuple[int]]
          The padding of convolution on both sides of inputs before convolution.
  
      dilation : Optional[int, Tuple[int]]
          Specifies the dilation rate to be used for dilated convolution.
  
      groups : Optional[int]
          Number of groups for grouped convolution.
  
      channels : Optional[int]
          Number of output channels of this convolution.
  
      kernel_size : Optional[int, Tuple[int]]
          The spatial of the convolution kernel.
  
      data_layout : Optional[str]
          Layout of the input.
  
      kernel_layout : Optional[str]
          Layout of the weight.
  
      out_layout : Optional[str]
          Layout of the output, by default, out_layout is the same as data_layout
  
      out_dtype : Optional[str]
          Specifies the output data type for mixed precision conv2d.
  
      Returns
      -------
      result : tvm.relay.Expr
          The computed result.
      """
      if isinstance(kernel_size, int):
          kernel_size = (kernel_size, kernel_size)
      if isinstance(strides, int):
          strides = (strides, strides)
      if isinstance(dilation, int):
          dilation = (dilation, dilation)
      # TODO enforce 4-way padding in topi/nn/conv2d after #4644 merged
      # convert 2-way padding to 4-way padding
      padding = get_pad_tuple2d(padding)
      return _make.conv2d(
          data,
          weight,
          strides,
          padding,
          dilation,
          groups,
          channels,
          kernel_size,
          data_layout,
          kernel_layout,
          out_layout,
          out_dtype,
      )
  
  
  def conv3d(
      data,
      weight,
      strides=(1, 1, 1),
      padding=(0, 0, 0),
      dilation=(1, 1, 1),
      groups=1,
      channels=None,
      kernel_size=None,
      data_layout="NCDHW",
      kernel_layout="OIDHW",
      out_layout="",
      out_dtype="",
  ):
      r"""3D 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 `NCDHW`
      and kernel_layout is `OIDHW`, conv3d takes in
      a data Tensor with shape `(batch_size, in_channels, depth, height, width)`,
      and a weight Tensor with shape `(channels, in_channels, kernel_size[0], kernel_size[1],
      kernel_size[2])` to produce an output Tensor with the following rule:
  
      .. math::
  
          \mbox{out}[b, c, z, y, x] = \sum_{dz, dy, dx, k}
             \mbox{data}[b, k, \mbox{strides}[0] * z  + dz, \mbox{strides}[1] * y  + dy,
             \mbox{strides}[2] * x + dx] * \mbox{weight}[c, k, dz, 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
      (`NCDHW` for data and `OIDHW` for weight), perform the computation,
      then convert to the out_layout.
  
  
      Parameters
      ----------
      data : tvm.relay.Expr
          The input data to the operator.
  
      weight : tvm.relay.Expr
          The weight expressions.
  
      strides : Optional[Tuple[int]]
          The strides of convolution.
  
      padding : Optional[int, Tuple[int]]
          The padding of convolution on both sides of inputs before convolution.
  
      dilation : Optional[int, Tuple[int]]
          Specifies the dilation rate to be used for dilated convolution.
  
      groups : Optional[int]
          Number of groups for grouped convolution.
  
      channels : Optional[int]
          Number of output channels of this convolution.
  
      kernel_size : Optional[int, Tuple[int]]
          The spatial of the convolution kernel.
  
      data_layout : Optional[str]
          Layout of the input.
  
      kernel_layout : Optional[str]
          Layout of the weight.
  
      out_layout : Optional[str]
          Layout of the output, by default, out_layout is the same as data_layout
  
      out_dtype : Optional[str]
          Specifies the output data type for mixed precision conv2d.
  
      Returns
      -------
      result : tvm.relay.Expr
          The computed result.
      """
      if isinstance(kernel_size, int):
          kernel_size = (kernel_size, kernel_size, kernel_size)
      if isinstance(strides, int):
          strides = (strides, strides, strides)
      if isinstance(dilation, int):
          dilation = (dilation, dilation, dilation)
      padding = get_pad_tuple3d(padding)
      return _make.conv3d(
          data,
          weight,
          strides,
          padding,
          dilation,
          groups,
          channels,
          kernel_size,
          data_layout,
          kernel_layout,
          out_layout,
          out_dtype,
      )
  
  
  def contrib_conv3d_winograd_without_weight_transform(
      data,
      weight,
      tile_size,
      strides=(1, 1, 1),
      padding=(0, 0, 0),
      dilation=(1, 1, 1),
      groups=1,
      channels=None,
      kernel_size=None,
      data_layout="NCDHW",
      kernel_layout="OIDHW",
      out_layout="",
      out_dtype="",
  ):
      r"""3D convolution with winograd algorithm.
  
      The basic parameters are the same as the ones in vanilla conv3d.
      It assumes the weight is pre-transformed by nn.contrib_conv3d_winograd_weight_transform
  
      Parameters
      ----------
      data : tvm.relay.Expr
          The input data to the operator.
  
      weight : tvm.relay.Expr
          The weight expressions.
  
      tile_size : int
          The Tile size of winograd. E.g. 2 for F(2x2x2, 3x3x3) and 4 for F(4x4x4, 3x3x3)
  
      strides : tuple of int, optional
          The strides of convolution.
  
      padding : tuple of int, optional
          The padding of convolution on both sides of inputs before convolution.
  
      dilation : tuple of int, optional
          Specifies the dilation rate to be used for dilated convolution.
  
      groups : int, optional
          Number of groups for grouped convolution.
  
      channels : int, optional
          Number of output channels of this convolution.
  
      kernel_size : tuple of int, optional
          The spatial of the convolution kernel.
  
      data_layout : str, optional
          Layout of the input.
  
      kernel_layout : str, optional
          Layout of the weight.
  
      out_layout : str, optional
          Layout of the output, by default, out_layout is the same as data_layout
  
      out_dtype : str, optional
          Specifies the output data type for mixed precision conv2d.
  
      Returns
      -------
      result : tvm.relay.Expr
          The computed result.
      """
      # convert 3-way padding to 6-way padding
      padding = get_pad_tuple3d(padding)
      return _make.contrib_conv3d_winograd_without_weight_transform(
          data,
          weight,
          tile_size,
          strides,
          padding,
          dilation,
          groups,
          channels,
          kernel_size,
          data_layout,
          kernel_layout,
          out_layout,
          out_dtype,
      )
  
  
  def conv3d_transpose(
      data,
      weight,
      strides=(1, 1, 1),
      padding=(0, 0, 0),
      dilation=(1, 1, 1),
      groups=1,
      channels=None,
      kernel_size=None,
      data_layout="NCDHW",
      kernel_layout="OIDHW",
      out_layout="",
      output_padding=(0, 0, 0),
      out_dtype="",
  ):
      r"""3D transpose convolution.
  
      Parameters
      ----------
      data : tvm.relay.Expr
          The input data to the operator.
  
      weight : tvm.relay.Expr
          The weight expressions.
  
      strides : Optional[Tuple[int]]
          The strides of convolution.
  
      padding : Optional[int, Tuple[int]]
          The padding of convolution on both sides of inputs before convolution.
  
      dilation : Optional[int, Tuple[int]]
          Specifies the dilation rate to be used for dilated convolution.
  
      groups : Optional[int]
          Number of groups for grouped convolution.
  
      channels : Optional[int]
          Number of output channels of this convolution.
  
      kernel_size : Optional[int, Tuple[int]]
          The spatial of the convolution kernel.
  
      data_layout : Optional[str]
          Layout of the input.
  
      kernel_layout : Optional[str]
          Layout of the weight.
  
      out_layout : Optional[str]
          Layout of the output, by default, out_layout is the same as data_layout
  
      out_dtype : Optional[str]
          Specifies the output data type for mixed precision conv3d.
  
      Returns
      -------
      result : tvm.relay.Expr
          The computed result.
      """
  
      if isinstance(kernel_size, int):
          kernel_size = (kernel_size, kernel_size, kernel_size)
      if isinstance(strides, int):
          strides = (strides, strides, strides)
      if isinstance(dilation, int):
          dilation = (dilation, dilation, dilation)
      padding = get_pad_tuple3d(padding)
  
      return _make.conv3d_transpose(
          data,
          weight,
          strides,
          padding,
          dilation,
          groups,
          channels,
          kernel_size,
          data_layout,
          kernel_layout,
          out_layout,
          output_padding,
          out_dtype,
      )
  
  
  def conv2d_transpose(
      data,
      weight,
      strides=(1, 1),
      padding=(0, 0),
      dilation=(1, 1),
      groups=1,
      channels=None,
      kernel_size=None,
      data_layout="NCHW",
      kernel_layout="OIHW",
      out_layout="",
      output_padding=(0, 0),
      out_dtype="",
  ):
      """Two dimensional transposed convolution operator.
  
      Parameters
      ----------
      data : tvm.relay.Expr
          The input data to the operator.
  
      weight : tvm.relay.Expr
          The weight expressions.
  
      strides : Tuple[int], optional
          The strides of convolution.
  
      padding : Tuple[int], optional
          The padding of convolution on both sides of inputs.
  
      dilation : Tuple[int], optional
          Specifies the dilation rate to be used for dilated convolution.
  
      channels : int, optional
          Number of output channels of this convolution.
  
      kernel_size : tuple of int, optional
          The spatial of the convolution kernel.
  
      groups : int, optional
          Number of groups for grouped convolution.
  
      data_layout : str, optional
          Layout of the input.
  
      kernel_layout : str, optional
          Layout of the weight.
  
      out_layout : Optional[str]
          Layout of the output, by default, out_layout is the same as data_layout
  
      output_padding : Tuple[int], optional
          Used to disambiguate the output shape.
  
      out_dtype : str, optional
          Specifies the output data type for mixed precision conv2d.
  
      Returns
      -------
      result : tvm.relay.Expr
          The computed result.
      """
      # convert 2-way padding to 4-way padding
      padding = get_pad_tuple2d(padding)
      return _make.conv2d_transpose(
          data,
          weight,
          strides,
          padding,
          dilation,
          groups,
          channels,
          kernel_size,
          data_layout,
          kernel_layout,
          out_layout,
          output_padding,
          out_dtype,
      )
  
  
  def conv1d_transpose(
      data,
      weight,
      strides=(1,),
      padding=(0,),
      dilation=(1,),
      groups=1,
      channels=None,
      kernel_size=None,
      data_layout="NCW",
      kernel_layout="OIW",
      out_layout="",
      output_padding=(0,),
      out_dtype="",
  ):
      """One dimensional transposed convolution operator.
  
      Parameters
      ----------
      data : tvm.relay.Expr
          The input data to the operator.
  
      weight : tvm.relay.Expr
          The weight expressions.
  
      strides : Tuple[int], optional
          The strides of convolution.
  
      padding : Tuple[int], optional
          The padding of convolution on both sides of inputs.
  
      dilation : Tuple[int], optional
          Specifies the dilation rate to be used for dilated convolution.
  
      channels : int, optional
          Number of output channels of this convolution.
  
      kernel_size : tuple of int, optional
          The spatial of the convolution kernel.
  
      groups : int, optional
          Number of groups for grouped convolution.
  
      data_layout : str, optional
          Layout of the input.
  
      kernel_layout : str, optional
          Layout of the weight.
  
      out_layout : Optional[str]
          Layout of the output, by default, out_layout is the same as data_layout
  
      output_padding : Tuple[int], optional
          Used to disambiguate the output shape.
  
      out_dtype : str, optional
          Specifies the output data type for mixed precision conv2d.
  
      Returns
      -------
      result : tvm.relay.Expr
          The computed result.
      """
      return _make.conv1d_transpose(
          data,
          weight,
          strides,
          padding,
          dilation,
          groups,
          channels,
          kernel_size,
          data_layout,
          kernel_layout,
          out_layout,
          output_padding,
          out_dtype,
      )
  
  
  def softmax(data, axis=-1):
      r"""Computes softmax.
  
      .. math:: \text{softmax}(x)_i = \frac{exp(x_i)}{\sum_j exp(x_j)}
  
      .. note::
          This operator can be optimized away for inference.
  
      Parameters
      ----------
      data: tvm.relay.Expr
          The input data to the operator.
  
      axis: int, optional
          The axis to sum over when computing softmax
  
      Returns
      -------
      result : tvm.relay.Expr
          The computed result.
      """
      return _make.softmax(data, axis)
  
  
  def fast_softmax(data, axis=-1):
      r"""Computes softmax.
      Use approximation to compute exponent for faster speed.
  
      .. math:: \text{softmax}(x)_i = \frac{exp(x_i)}{\sum_j exp(x_j)}
      .. note::
          This operator can be optimized away for inference.
  
      Parameters
      ----------
      data: tvm.relay.Expr
          The input data to the operator.
      axis: int, optional
          The axis to sum over when computing softmax
  
      Returns
      -------
      result : tvm.relay.Expr
          The computed result.
      """
      return _make.fast_softmax(data, axis)
  
  
  def log_softmax(data, axis=-1):
      r"""Computes log softmax.
  
      .. math::
  
          \text{log_softmax}(x)_i = \log \frac{exp(x_i)}{\sum_j exp(x_j)}
  
      .. note::
          This operator can be optimized away for inference.
  
      Parameters
      ----------
      data: tvm.relay.Expr
          The input data to the operator.
  
      axis: int, optional
          The axis to sum over when computing log softmax
  
      Returns
      -------
      result : tvm.relay.Expr
          The computed result.
      """
      return _make.log_softmax(data, axis)
  
  
  def max_pool1d(data, pool_size=(1,), strides=(1,), padding=(0,), layout="NCW", ceil_mode=False):
      r"""1D maximum pooling operator.
  
      This operator takes data as input and does 1D max value calculation
      with in pool_size sized window by striding defined by stride.
  
      In the default case, where the data_layout is `NCW`
      a data Tensor with shape `(batch_size, channels, width)`,
      to produce an output Tensor.
  
      The ceil_mode is used to take ceil or floor while computing out shape.
      count_include_pad indicates including or excluding padded input values in computation.
      This operator accepts data layout specification.
  
      Parameters
      ----------
      data : tvm.relay.Expr
          The input data to the operator.
  
      pool_size : int or tuple of int, optional
          The size of window for pooling.
  
      strides : int or tuple of int, optional
          The strides of pooling.
  
      padding : int or tuple of int, optional
          The padding for pooling.
  
      layout : str, optional
          Layout of the input.
  
      ceil_mode : bool, optional
          To enable or disable ceil while pooling.
  
      Returns
      -------
      result : tvm.relay.Expr
          The computed result.
      """
      if isinstance(pool_size, int):
          pool_size = (pool_size,)
      if isinstance(strides, int):
          strides = (strides,)
      padding = get_pad_tuple1d(padding)
      return _make.max_pool1d(data, pool_size, strides, padding, layout, ceil_mode)
  
  
  def max_pool2d(
      data, pool_size=(1, 1), strides=(1, 1), padding=(0, 0), layout="NCHW", ceil_mode=False
  ):
      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 : tvm.relay.Expr
          The input data to the operator.
  
      pool_size : int or tuple of int, optional
          The size of window for pooling.
  
      strides : tuple of int, optional
          The strides of pooling.
  
      padding : tuple of int, optional
          The padding for pooling.
  
      layout : str, optional
          Layout of the input.
  
      ceil_mode : bool, optional
          To enable or disable ceil while pooling.
  
      Returns
      -------
      result : tvm.relay.Expr
          The computed result.
      """
      if isinstance(pool_size, int):
          pool_size = (pool_size, pool_size)
      if isinstance(strides, int):
          strides = (strides, strides)
      padding = get_pad_tuple2d(padding)
      return _make.max_pool2d(data, pool_size, strides, padding, layout, ceil_mode)
  
  
  def max_pool3d(
      data, pool_size=(1, 1, 1), strides=(1, 1, 1), padding=(0, 0, 0), layout="NCDHW", ceil_mode=False
  ):
      r"""3D maximum pooling operator.
  
      This operator takes data as input and does 3D max value calculation
      with in pool_size sized window by striding defined by stride.
  
  
      In the default case, where the data_layout is `NCDHW`
      a data Tensor with shape `(batch_size, channels, depth, height, width)`,
      to produce an output Tensor.
  
      The ceil_mode is used to take ceil or floor while computing out shape.
      count_include_pad indicates including or excluding padded input values in computation.
      This operator accepts data layout specification.
  
      Parameters
      ----------
      data : tvm.relay.Expr
          The input data to the operator.
  
      pool_size : int or tuple of int, optional
          The size of window for pooling.
  
      strides : tuple of int, optional
          The strides of pooling.
  
      padding : tuple of int, optional
          The padding for pooling.
  
      layout : str, optional
          Layout of the input.
  
      ceil_mode : bool, optional
          To enable or disable ceil while pooling.
  
      Returns
      -------
      result : tvm.relay.Expr
          The computed result.
      """
      if isinstance(pool_size, int):
          pool_size = (pool_size, pool_size, pool_size)
      if isinstance(strides, int):
          strides = (strides, strides, strides)
      padding = get_pad_tuple3d(padding)
      return _make.max_pool3d(data, pool_size, strides, padding, layout, ceil_mode)
  
  
  def avg_pool1d(
      data,
      pool_size=(1,),
      strides=(1,),
      padding=(0,),
      layout="NCW",
      ceil_mode=False,
      count_include_pad=False,
  ):
      r"""1D average pooling operator.
  
      This operator takes data as input and does 1D average value calculation
      with in pool_size sized window by striding defined by stride
  
      In the default case, where the data_layout is `NCW`
      a data Tensor with shape `(batch_size, channels, width)`,
      to produce an output Tensor.
  
      The ceil_mode is used to take ceil or floor while computing out shape.
      count_include_pad indicates including or excluding padded input values in computation.
      This operator accepts data layout specification.
  
      Parameters
      ----------
      data : tvm.relay.Expr
          The input data to the operator.
  
      pool_size : int or tuple of int, optional
          The size of window for pooling.
  
      strides : int or tuple of int, optional
          The strides of pooling.
  
      padding : int or tuple of int, optional
          The padding for pooling.
  
      layout : str, optional
          Layout of the input.
  
      ceil_mode : bool, optional
          To enable or disable ceil while pooling.
  
      count_include_pad : bool, optional
          To include padding to compute the average.
  
      Returns
      -------
      result : tvm.relay.Expr
          The computed result.
      """
      if isinstance(pool_size, int):
          pool_size = (pool_size,)
      if isinstance(strides, int):
          strides = (strides,)
      padding = get_pad_tuple1d(padding)
      return _make.avg_pool1d(data, pool_size, strides, padding, layout, ceil_mode, count_include_pad)
  
  
  def avg_pool2d(
      data,
      pool_size=(1, 1),
      strides=(1, 1),
      padding=(0, 0),
      layout="NCHW",
      ceil_mode=False,
      count_include_pad=False,
  ):
      r"""2D average pooling operator.
  
      This operator takes data as input and does 2D average 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), pool_size (kh, kw)
  
      .. math::
  
          \mbox{out}(b, c, y, x)  = \frac{1}{kh * kw} \sum_{m=0}^{kh-1} \sum_{n=0}^{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.
      count_include_pad indicates including or excluding padded input values in computation.
      This operator accepts data layout specification.
  
      Parameters
      ----------
      data : tvm.relay.Expr
          The input data to the operator.
  
      pool_size : int or tuple of int, optional
          The size of window for pooling.
  
      strides : tuple of int, optional
          The strides of pooling.
  
      padding : tuple of int, optional
          The padding for pooling.
  
      layout : str, optional
          Layout of the input.
  
      ceil_mode : bool, optional
          To enable or disable ceil while pooling.
  
      count_include_pad : bool, optional
          To include padding to compute the average.
  
      Returns
      -------
      result : tvm.relay.Expr
          The computed result.
      """
      if isinstance(pool_size, int):
          pool_size = (pool_size, pool_size)
      if isinstance(strides, int):
          strides = (strides, strides)
      padding = get_pad_tuple2d(padding)
      return _make.avg_pool2d(data, pool_size, strides, padding, layout, ceil_mode, count_include_pad)
  
  
  def avg_pool3d(
      data,
      pool_size=(1, 1, 1),
      strides=(1, 1, 1),
      padding=(0, 0, 0),
      layout="NCDHW",
      ceil_mode=False,
      count_include_pad=False,
  ):
      r"""3D average pooling operator.
  
      This operator takes data as input and does 3D average value calculation
      with in pool_size sized window by striding defined by stride
  
  
      In the default case, where the data_layout is `NCDHW`
      a data Tensor with shape `(batch_size, channels, depth, height, width)`,
      to produce an output Tensor.
  
      The ceil_mode is used to take ceil or floor while computing out shape.
      count_include_pad indicates including or excluding padded input values in computation.
      This operator accepts data layout specification.
  
      Parameters
      ----------
      data : tvm.relay.Expr
          The input data to the operator.
  
      pool_size : int or tuple of int, optional
          The size of window for pooling.
  
      strides : tuple of int, optional
          The strides of pooling.
  
      padding : tuple of int, optional
          The padding for pooling.
  
      layout : str, optional
          Layout of the input.
  
      ceil_mode : bool, optional
          To enable or disable ceil while pooling.
  
      count_include_pad : bool, optional
          To include padding to compute the average.
  
      Returns
      -------
      result : tvm.relay.Expr
          The computed result.
      """
      if isinstance(pool_size, int):
          pool_size = (pool_size, pool_size, pool_size)
      if isinstance(strides, int):
          strides = (strides, strides, strides)
      padding = get_pad_tuple3d(padding)
      return _make.avg_pool3d(data, pool_size, strides, padding, layout, ceil_mode, count_include_pad)
  
  
  def max_pool2d_grad(
      out_grad, data, pool_size=(1, 1), strides=(1, 1), padding=(0, 0), layout="NCHW", ceil_mode=False
  ):
      r"""Gradient of 2D maximum pooling operator.
  
      This operator takes out_grad and data as input and calculates gradient of max_pool2d.
  
      Parameters
      ----------
      out_grad : tvm.relay.Expr
          The output gradient
  
      data : tvm.relay.Expr
          The input data to the operator.
  
      pool_size : int or tuple of int, optional
          The size of window for pooling.
  
      strides : tuple of int, optional
          The strides of pooling.
  
      padding : tuple of int, optional
          The padding for pooling.
  
      layout : str, optional
          Layout of the input.
  
      ceil_mode : bool, optional
          To enable or disable ceil while pooling.
  
      Returns
      -------
      result : tvm.relay.Expr
          The computed result.
      """
      return _make.max_pool2d_grad(out_grad, data, pool_size, strides, padding, layout, ceil_mode)
  
  
  def avg_pool2d_grad(
      out_grad,
      data,
      pool_size=(1, 1),
      strides=(1, 1),
      padding=(0, 0),
      layout="NCHW",
      ceil_mode=False,
      count_include_pad=False,
  ):
      r"""Gradient of 2D average pooling operator.
  
      This operator takes out_grad and data as input and calculates gradient of avg_pool2d.
  
      Parameters
      ----------
      out_grad : tvm.relay.Expr
          The output gradient
  
      data : tvm.relay.Expr
          The input data to the operator.
  
      pool_size : int or tuple of int, optional
          The size of window for pooling.
  
      strides : tuple of int, optional
          The strides of pooling.
  
      padding : tuple of int, optional
          The padding for pooling.
  
      layout : str, optional
          Layout of the input.
  
      ceil_mode : bool, optional
          To enable or disable ceil while pooling.
  
      count_include_pad : bool, optional
          To include padding to compute the average.
  
      Returns
      -------
      result : tvm.relay.Expr
          The computed result.
      """
      return _make.avg_pool2d_grad(
          out_grad, data, pool_size, strides, padding, layout, ceil_mode, count_include_pad
      )
  
  
  def global_max_pool2d(data, layout="NCHW"):
      r"""2D global maximum pooling operator.
  
      This operator takes data as input and does 2D max 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 the following rule:
  
      with data of shape (b, c, h, w)
  
      .. math::
  
          \mbox{out}(b, c, 1, 1)  = \max_{m=0, \ldots, h} \max_{n=0, \ldots, w}
               \mbox{data}(b, c, m, n)
  
      Parameters
      ----------
      data : tvm.relay.Expr
          The input data to the operator.
  
      layout : str, optional
          Layout of the input.
  
      Returns
      -------
      result : tvm.relay.Expr
          The computed result.
      """
      return _make.global_max_pool2d(data, layout)
  
  
  def global_avg_pool2d(data, layout="NCHW"):
      r"""2D global average pooling operator.
  
      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 the following rule:
  
      with data of shape (b, c, h, w)
  
      .. math::
  
          \mbox{out}(b, c, 1, 1)  = \frac{1}{h * w} \sum_{m=0}^{h-1} \sum_{n=0}^{w-1}
               \mbox{data}(b, c, m, n)
  
      Parameters
      ----------
      data : tvm.relay.Expr
          The input data to the operator.
  
      layout : str, optional
          Layout of the input.
  
      Returns
      -------
      result : tvm.relay.Expr
          The computed result.
      """
      return _make.global_avg_pool2d(data, layout)
  
  
  def upsampling(
      data, scale_h=1, scale_w=1, layout="NCHW", method="nearest_neighbor", align_corners=False
  ):
      """Upsampling.
  
      This operator takes data as input and does 2D scaling to the given scale factor.
      In the default case, where the data_layout is `NCHW`
      with data of shape (n, c, h, w)
      out will have a shape (n, c, h*scale_h, w*scale_w)
  
      method indicates the algorithm to be used while calculating the out value
      and method can be one of ("bilinear", "nearest_neighbor", "bicubic")
  
      Parameters
      ----------
      data : tvm.relay.Expr
          The input data to the operator.
  
      scale_h : tvm.relay.Expr or int or float
          The scale factor for height upsampling.
  
      scale_w : tvm.relay.Expr or int or float
          The scale factor for width upsampling.
  
      layout : str, optional
          Layout of the input.
  
      method : str, optional
          Scale method to used [nearest_neighbor, bilinear, bicubic].
  
      align_corners : bool, optional
          Whether to keep corners in proper place.
  
      Returns
      -------
      result : tvm.relay.Expr
          The computed result.
      """
      if isinstance(scale_h, Expr) or isinstance(scale_w, Expr):
          if not isinstance(scale_h, Expr):
              scale_h = const(scale_h, "float64")
          if not isinstance(scale_w, Expr):
              scale_w = const(scale_w, "float64")
          return _dyn_make.upsampling(data, scale_h, scale_w, layout, method, align_corners)
      return _make.upsampling(data, scale_h, scale_w, layout, method, align_corners)
  
  
  def upsampling3d(
      data,
      scale_d=1,
      scale_h=1,
      scale_w=1,
      layout="NCDHW",
      method="nearest_neighbor",
      coordinate_transformation_mode="half_pixel",
  ):
      """3D Upsampling.
  
      This operator takes data as input and does 3D scaling to the given scale factor.
      In the default case, where the data_layout is `NCDHW`
      with data of shape (n, c, d, h, w)
      out will have a shape (n, c, d*scale_d, h*scale_h, w*scale_w)
  
      method indicates the algorithm to be used while calculating the out value
      and method can be one of ("trilinear", "nearest_neighbor")
  
      Parameters
      ----------
      data : tvm.relay.Expr
          The input data to the operator.
  
      scale_d : tvm.relay.Expr
          The scale factor for depth upsampling.
  
      scale_h : tvm.relay.Expr
          The scale factor for height upsampling.
  
      scale_w : tvm.relay.Expr
          The scale factor for width upsampling.
  
      layout : str, optional
          Layout of the input.
  
      method : str, optional
          Scale method to used [nearest_neighbor, trilinear].
  
      coordinate_transformation_mode: string, optional
          Describes how to transform the coordinate in the resized tensor
          to the coordinate in the original tensor.
          Refer to the ONNX Resize operator specification for details.
          Available options are "half_pixel", "align_corners" and "asymmetric".
  
      Returns
      -------
      result : tvm.relay.Expr
          The computed result.
      """
      if isinstance(scale_d, Expr) or isinstance(scale_h, Expr) or isinstance(scale_w, Expr):
          if not isinstance(scale_d, Expr):
              scale_d = const(scale_d, "float64")
          if not isinstance(scale_h, Expr):
              scale_h = const(scale_h, "float64")
          if not isinstance(scale_w, Expr):
              scale_w = const(scale_w, "float64")
          return _dyn_make.upsampling3d(
              data, scale_d, scale_h, scale_w, layout, method, coordinate_transformation_mode
          )
      return _make.upsampling3d(
          data, scale_d, scale_h, scale_w, layout, method, coordinate_transformation_mode
      )
  
  
  def batch_flatten(data):
      """BatchFlatten.
  
      This operator flattens all the dimensions except for the batch dimension.
      which results a 2D output.
  
      For data with shape ``(d1, d2, ..., dk)``
      batch_flatten(data) returns reshaped output of shape ``(d1, d2*...*dk)``.
  
  
      Parameters
      ----------
      data : tvm.relay.Expr
          The input data to the operator.
  
      Returns
      -------
      result : tvm.relay.Expr
          The Flattened result.
      """
      return _make.batch_flatten(data)
  
  
  def bias_add(data, bias, axis=1):
      """add_bias operator.
  
      Add 1D bias to the axis of data.
      This function is a special case of add which allows
      inference of shape of the bias from data.
  
      Parameters
      ----------
      data : tvm.relay.Expr
          The input data to the operator.
  
      bias : tvm.relay.Expr
          The bias to be added.
  
      axis : int, optional
          The axis to add the bias.
  
      Returns
      -------
      result : tvm.relay.Expr
          The final result.
      """
      return _make.bias_add(data, bias, axis)
  
  
  def dense(data, weight, units=None, out_dtype=""):
      """Dense operator.
      Applies a linear transformation
  
      .. math::
  
      `Y = X * W^T`
  
      Parameters
      ----------
      data : tvm.relay.Expr
          The input data to the operator,
          of shape `(d_1, d_2, ..., d_n, units_in)`.
  
      weight : tvm.relay.Expr
          The weight expressions, 2-D matrix,
          of shape `(units, units_in)`.
  
      units : int, optional
          Number of hidden units of the dense transformation.
  
      out_dtype : str, optional
          Specifies the output data type for mixed precision dense,
          of shape `(d_1, d_2, ..., d_n, units)`.
  
      Returns
      -------
      result : tvm.relay.Expr
          The computed result.
      """
      return _make.dense(data, weight, units, out_dtype)
  
  
  def fifo_buffer(data, buffer, axis):
      """FIFO buffer to enable computation reuse in CNNs with sliding indow input
  
      Compute equivalent of
  
      .. code-block:: python
  
          concat(buffer, data, axis=axis)
          .slice_axis(axis=axis,
                      begin=data.shape[axis],
                      end=data.shape[axis]+buffer.shape[axis])
  
      Useful for
  
      * Encoding explicit re-use of computation in convolution ops operated on a sliding window input
      * Implementing a FIFO queue to cache intermediate results, e.g. as in Fast WaveNet.
  
      Parameters
      ----------
      data : tvm.relay.Expr
          The input data
      buffer : tvm.relay.Expr
          Previous value of the FIFO buffer
      axis : int
          Specify which axis should be used for buffering
  
      Returns
      -------
      result : tvm.relay.Expr
          Updated value for the buffer
      """
      return _make.fifo_buffer(data, buffer, axis)
  
  
  def relu(data):
      """Rectified linear unit.
  
      .. math::
         out = max(x, 0)
  
      Parameters
      ----------
      data : tvm.relay.Expr
          The input data
  
      Returns
      -------
      result : tvm.relay.Expr
          The computed result.
      """
      return _make.relu(data)
  
  
  def leaky_relu(data, alpha=0.01):
      """This operator takes data as input and does Leaky version
      of a Rectified Linear Unit.
  
      .. math::
  
          `y = x > 0 ? x : alpha * x`
  
      Parameters
      ----------
      data : tvm.relay.Expr
          The input data to the operator.
  
      alpha : float
          Slope coefficient for the negative half axis.
  
      Returns
      -------
      result : tvm.relay.Expr
          The computed result.
      """
      return _make.leaky_relu(data, alpha)
  
  
  def prelu(data, alpha, axis=1):
      """This operator takes data as input and does Leaky version
      of a Rectified Linear Unit.
  
      .. math::
  
          y = x > 0 ? x : alpha * x
  
      Parameters
      ----------
      data : tvm.relay.Expr
          The input data to the operator.
  
      alpha : tvm.relay.Expr
          Slope coefficient for the negative half axis.
  
      axis : int, optional
          Specify which shape axis the channel is specified.
  
      Returns
      -------
      result : tvm.relay.Expr
          The computed result.
      """
      return _make.prelu(data, alpha, axis)
  
  
  def pad(data, pad_width, pad_value=0, pad_mode="constant"):
      r"""Padding
  
      This operator takes in a tensor and pads each axis by the specified
      widths using the specified value.
  
      Parameters
      ----------
      data: tvm.relay.Expr
          The input data to the operator
      pad_width: tuple of <tuple of <int>>, or tvm.relay.Expr, required
          Number of values padded to the edges of each axis, in the format
          of ((before_1, after_1), ..., (before_N, after_N))
      pad_value: float, or tvm.relay.Expr, optional, default=0
          The value used for padding
      pad_mode: 'constant', 'edge', 'reflect'
          'constant' pads with constant_value pad_value
          'edge' pads using the edge values of the input array
          'reflect' pads by reflecting values with respect to the edge
      Returns
      -------
      result : tvm.relay.Expr
          The computed result.
      """
      if isinstance(pad_width, Expr) or (isinstance(pad_value, Expr)):
          if not isinstance(pad_width, Expr):
              pad_width = const(list(pad_width))
          if not isinstance(pad_value, Expr):
              pad_value = const(pad_value)
          return _dyn_make.pad(data, pad_width, pad_value, pad_mode)
      return _make.pad(data, pad_width, pad_value, pad_mode)
  
  
  def dilate(data, strides, dilation_value=0.0):
      """Dilate data with given dilation value (0 by default).
  
      Parameters
      ----------
      data : tvm.relay.Expr
          n-D, can be any layout.
  
      strides : tuple of <int>
          Dilation stride on each dimension, 1 means no dilation.
  
      dilation_value : int/float, optional
          Value used to dilate the input.
  
      Returns
      -------
      Output : tvm.relay.Expr
          The computed result
      """
      return _make.dilate(data, strides, dilation_value)
  
  
  def mirror_pad(data, pad_width, mode="SYMMETRIC"):
      r"""MirrorPadding
  
      This operator takes in a tensor and pads each axis by the specified
      widths using mirroring of the border pixels.
  
      Parameters
      ----------
      data: tvm.relay.Expr
          The input data to the operator
      pad_width: tuple of <tuple of <int>>, required
          Number of values padded to the edges of each axis, in the format
          of ((before_1, after_1), ..., (before_N, after_N))
      mode: string, optional, default='SYMMETRIC'
          What type of mirroring to use, must be SYMMETRIC or REFLECT.
  
      Returns
      -------
      result : tvm.relay.Expr
          The computed result.
      """
      return _make.mirror_pad(data, pad_width, mode)
  
  
  def lrn(data, size=5, axis=1, bias=2, alpha=0.00001, beta=0.75):
      """This operator takes data as input and does local response normalization.
  
      Normalize the input in a local region across or within feature maps.
      Each input value is divided by (data / (bias + (alpha * sum_data ^2 /size))^beta)
      where n is the size of each local region, and the sum is taken over the region
      centered at that value (zero padding is added where necessary).
  
      .. math::
          (data / (bias + (alpha * sum_data ^2 /size))^beta)
  
      Parameters
      ----------
      data : tvm.relay.Expr
          The input data to the operator.
  
      size : int, optional
          The size of the local region to be considered for normalization.
  
      axis : int, optional
          Input data layout channel axis. Default value is 1 for NCHW format
  
      bias : float, optional
          The offset parameter to avoid dividing by 0.
  
      alpha : float, optional
          The scaling parameter.
  
      beta : float, optional
          The exponent parameter.
  
      Returns
      -------
      result : tvm.relay.Expr
          The computed result.
      """
      return _make.lrn(data, size, axis, alpha, beta, bias)
  
  
  def l2_normalize(data, eps, axis=None):
      """Perform L2 normalization on the input data
  
      .. math::
          y(i, j) = x(i, j) / sqrt(max(sum(x^2), eps))
  
      Parameters
      ----------
      data : tvm.relay.Expr
          The input data to the operator.
  
      eps : float
          epsilon value
  
      axis : list of int, optional
          axis over the normalization applied
  
      Returns
      -------
      result : tvm.relay.Expr
          The computed result.
      """
      return _make.l2_normalize(data, eps, axis)
  
  
  def dropout(data, rate=0.5):
      """Applies the dropout operation to the input array.
  
      During training, each element of the input is set to zero with
      probability ``p``. The whole array is rescaled by ``1/(1-p)``
      to keep the expected sum of the input unchanged.
  
      Parameters
      ----------
      data : tvm.relay.Expr
          The input data to the operator.
  
      rate : float, optional (default=0.5)
          The probability for an element to be reset to 0.
  
      Returns
      -------
      result : tvm.relay.Expr
          The result of dropout
      """
      return expr.TupleWrapper(dropout_raw(data, rate), 2)[0]
  
  
  def dropout_raw(data, rate=0.5):
      """Applies the dropout operation to the input array.
  
      During training, each element of the input is set to zero with
      probability ``p``. The whole array is rescaled by ``1/(1-p)``
      to keep the expected sum of the input unchanged.
  
      Parameters
      ----------
      data : tvm.relay.Expr
          The input data to the operator.
  
      rate : float, optional (default=0.5)
          The probability for an element to be reset to 0.
  
      Returns
      -------
      result : tvm.relay.Expr
          The result of dropout
      """
      return _make.dropout(data, rate)
  
  
  def batch_norm(
      data, gamma, beta, moving_mean, moving_var, axis=1, epsilon=1e-5, center=True, scale=True
  ):
      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 : tvm.relay.Expr
          Input to which batch_norm will be applied.
  
      gamma : tvm.relay.Expr
          The gamma scale factor.
  
      beta : tvm.relay.Expr
          The beta offset factor.
  
      moving_mean : tvm.relay.Expr
          Running mean of input,
  
      moving_var : tvm.relay.Expr
          Running variance of input.
  
      axis : int, optional, default=1
          Specify along which shape axis the channel is specified.
  
      epsilon : double, optional, default=1e-5
          Small float added to variance to avoid dividing by zero.
  
      center : boolean, optional, default=True
          If True, add offset of beta to normalized tensor, If False,
          beta is ignored.
  
      scale : boolean, optional, default=True
          If true, multiply by gamma. If False, gamma is not used.
          When the next layer is piecewise linear (also e.g. nn.relu),
          this can be disabled since the scaling will be done by the next layer.
  
      Returns
      -------
      result : relay.Tuple([tvm.relay.Expr, tvm.relay.Expr, tvm.relay.Expr])
          Tuple of normed data (same shape as input),
          new running mean (k-length vector),
          and new running variance (k-length vector)
      """
      result = _make.batch_norm(
          data, gamma, beta, moving_mean, moving_var, axis, epsilon, center, scale
      )
      return expr.TupleWrapper(result, 3)
  
  
  def instance_norm(data, gamma, beta, axis=1, epsilon=1e-5, center=True, scale=True):
      r"""
      Instance Normalization (Ulyanov and et al., 2016)
      Applies instance normalization to the n-dimensional input array.
  
      .. math::
  
          out = \frac{data - mean(data)}{\sqrt{var(data)+\epsilon}}
              * gamma + beta
  
      The instance normalization is similar to batch normalization, but unlike
      batch normalization, the mean and var are calculated per-dimension
      separately for each object(instance) in a mini-batch, not over a batch.
      And the same normalization is applied both at test and train time.
  
      Assume the input has size *k* on axis 1, then both ``gamma`` and ``beta``
      have shape *(k,)*.
  
      The parameter ``axis`` specifies which axis of the input shape denotes
      the 'channel'.  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 : tvm.relay.Expr
          Input to which instance_norm will be applied.
  
      gamma : tvm.relay.Expr
          The gamma scale factor.
  
      beta : tvm.relay.Expr
          The beta offset factor.
  
      axis : int, optional, default=1
          Specify along which shape axis the channel is specified.
  
      epsilon : double, optional, default=1e-5
          Small float added to variance to avoid dividing by zero.
  
      center : boolean, optional, default=True
          If True, add offset of beta to normalized tensor, If False,
          beta is ignored.
  
      scale : boolean, optional, default=True
          If True, multiply by gamma. If False, gamma is not used.
  
      Returns
      -------
      result : tvm.relay.Expr
          The normalized data.
  
      .. _`Instance Normalization: The Missing Ingredient for Fast Stylization`:
          https://arxiv.org/abs/1607.08022
      """
      return _make.instance_norm(data, gamma, beta, axis, epsilon, center, scale)
  
  
  def layer_norm(data, gamma, beta, axis=-1, epsilon=1e-5, center=True, scale=True):
      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 : tvm.relay.Expr
          Input to which layer_norm will be applied.
  
      gamma : tvm.relay.Expr
          The gamma scale factor.
  
      beta : tvm.relay.Expr
          The beta offset factor.
  
      axis : int, optional, default=-1
          The axis that should be normalized, typically the axis of the channels.
  
      epsilon : double, optional, default=1e-5
          Small float added to variance to avoid dividing by zero.
  
      center : boolean, optional, default=True
          If True, add offset of beta to normalized tensor, If False,
          beta is ignored.
  
      scale : boolean, optional, default=True
          If True, multiply by gamma. If False, gamma is not used.
  
      Returns
      -------
      result : tvm.relay.Expr
          The normalized data.
      """
      return _make.layer_norm(data, gamma, beta, axis, epsilon, center, scale)
  
  
  def group_norm(data, gamma, beta, num_groups, axis=1, epsilon=1e-5, center=True, scale=True):
      r"""
      Group normalization normalizes over group of channels for each training examples.
      We can say that, Group Norm is in between Instance Norm and Layer Norm. When we put
      all the channels into a single group, group normalization becomes Layer normalization.
      And, when we put each channel into different groups it becomes Instance normalization
  
      https://arxiv.org/pdf/1803.08494.pdf
  
      Applies group normalization to the n-dimensional input array by seperating the input channels
      into 'num_groups' groups, each containing 'num_channels / num_groups' channels.
      The mean and standard-deviation are calculated separately over the each group. gamma and
      beta are learnable per-channel affine transform parameter vectors of size num_channels.
  
      .. math::
  
          out = \frac{data - mean(data, axis)}{\sqrt{var(data, axis)+\epsilon}}
              * gamma + beta
  
      Unlike batch normalization, the mean and var are computed along a group of channels.
  
      If 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 : tvm.relay.Expr
          Input to which group_norm will be applied.
  
      gamma : tvm.relay.Expr
          The gamma scale factor.
  
      beta : tvm.relay.Expr
          The beta offset factor.
  
      num_groups : int
          The number of groups to separate the channels into.
  
      axis : int, optional, default=1
          The axis of the channels.
  
      epsilon : double, optional, default=1e-5
          Small float added to variance to avoid dividing by zero.
  
      center : boolean, optional, default=True
          If True, add offset of beta to normalized tensor, If False,
          beta is ignored.
  
      scale : boolean, optional, default=True
          If True, multiply by gamma. If False, gamma is not used.
  
      Returns
      -------
      result : tvm.relay.Expr
          The normalized data.
      """
      return _make.group_norm(data, gamma, beta, num_groups, axis, epsilon, center, scale)
  
  
  def batch_matmul(x, y):
      r"""
      Computes batch matrix multiplication of `x` and `y` when `x` and `y` are data
      in batch.
  
      .. math::
  
          \mbox{batch_matmul}(x, y)[i, :, :] = \mbox{matmul}(x[i, :, :], y[i, :, :]^T)
  
      Parameters
      ----------
      x : tvm.relay.Expr
          The first input.
  
      y : tvm.relay.Expr
          The second input.
  
      Returns
      -------
      result: tvm.relay.Expr
          The computed result.
      """
      return _make.batch_matmul(x, y)
  
  
  # pylint: disable=no-else-return,inconsistent-return-statements
  def sparse_dense(dense_mat, sparse_mat, sparse_lhs=False):
      r"""
      Computes the matrix multiplication of `dense_mat` and `sparse_mat`, where `dense_mat` is
      a dense matrix and `sparse_mat` is a sparse (either BSR or CSR) namedtuple with
      fields `data`, `indices`, and `indptr`.
  
      \if sparse_lhs=False:
          .. math::
  
              \mbox{sparse_dense}(dense_mat, sparse_mat)[m, n]
              = \mbox{matmul}(D, \mbox{as_dense}(S)^T)[m, n]
  
      \if sparse_lhs=True:
          .. math::
  
              \mbox{sparse_dense}(dense_mat, sparse_mat)[m, n]
              = \mbox{matmul}(\mbox{as_dense}(S), (D)^T)[m, n]
  
      where `as_dense` returns dense equivalent of the given S(sparse matrix)
      while performing matmul with given D(dense matrix).
  
      See
      https://docs.scipy.org/doc/scipy/reference/generated/scipy.sparse.csr_matrix.html
      and
      https://docs.scipy.org/doc/scipy-0.14.0/reference/generated/scipy.sparse.bsr_matrix.html
      for more detail on the sparse matrix representation.
  
      Parameters
      ----------
      dense_mat : tvm.relay.Expr
          The input dense matrix for the matrix multiplication
  
      sparse_mat : Union[namedtuple, Tuple[ndarray, ndarray, ndarray]].
          The input sparse matrix for the matrix multiplication.
  
      sparse_lhs : bool, optional
          Indicates whether lhs or rhs matrix is sparse. Default value is False.
  
      Returns
      -------
      result: tvm.relay.Expr
          The computed result.
      """
      if hasattr(sparse_mat, "indices"):
          return _make.sparse_dense(
              dense_mat, sparse_mat.data, sparse_mat.indices, sparse_mat.indptr, sparse_lhs
          )
      else:
          return _make.sparse_dense(
              dense_mat, sparse_mat[0], sparse_mat[1], sparse_mat[2], sparse_lhs
          )
  
  
  def sparse_transpose(x):
      r"""
      Computes the fast matrix transpose of x,
      where x is a sparse tensor in CSR format (represented as a namedtuple
      with fields `data`, `indices`, and `indptr`).
  
      ** Currently only support Square Matrices **
  
      .. math::
  
          \mbox{sparse_transpose}(x)[n, n] = (x^T)[n, n]
  
      Please refer to https://github.com/scipy/scipy/blob/v1.3.0/scipy/sparse/csr.py
      for the algorithm implemented in this operator.
  
      Parameters
      ----------
      x : Union[namedtuple, Tuple[ndarray, ndarray, ndarray]].
          The sparse weight matrix for the fast matrix transpose.
  
      Returns
      -------
      result : relay.Tuple([tvm.relay.Expr, tvm.relay.Expr, tvm.relay.Expr])
          Tuple of output sparse tensor (same shape and format as input),
          i.e. if CSR then output is in ([data, indices, indptr]) form
      """
      if hasattr(x, "indices"):
          return expr.TupleWrapper(_make.sparse_transpose(x.data, x.indices, x.indptr), 3)
      return expr.TupleWrapper(_make.sparse_transpose(x[0], x[1], x[2]), 3)
  
  
  def contrib_conv2d_winograd_without_weight_transform(
      data,
      weight,
      tile_size,
      strides=(1, 1),
      padding=(0, 0),
      dilation=(1, 1),
      groups=1,
      channels=None,
      kernel_size=None,
      data_layout="NCHW",
      kernel_layout="OIHW",
      out_layout="",
      out_dtype="",
  ):
      r"""2D convolution with winograd algorithm.
  
      The basic parameters are the same as the ones in vanilla conv2d.
      It assumes the weight is pre-transformed by nn.contrib_conv2d_winograd_weight_transform
  
      Parameters
      ----------
      data : tvm.relay.Expr
          The input data to the operator.
  
      weight : tvm.relay.Expr
          The weight expressions.
  
      tile_size : int
          The Tile size of winograd. E.g. 2 for F(2x2, 3x3) and 4 for F(4x4, 3x3)
  
      strides : tuple of int, optional
          The strides of convolution.
  
      padding : tuple of int, optional
          The padding of convolution on both sides of inputs before convolution.
  
      dilation : tuple of int, optional
          Specifies the dilation rate to be used for dilated convolution.
  
      groups : int, optional
          Number of groups for grouped convolution.
  
      channels : int, optional
          Number of output channels of this convolution.
  
      kernel_size : tuple of int, optional
          The spatial of the convolution kernel.
  
      data_layout : str, optional
          Layout of the input.
  
      kernel_layout : str, optional
          Layout of the weight.
  
      out_layout : str, optional
          Layout of the output, by default, out_layout is the same as data_layout
  
      out_dtype : str, optional
          Specifies the output data type for mixed precision conv2d.
  
      Returns
      -------
      result : tvm.relay.Expr
          The computed result.
      """
      # convert 2-way padding to 4-way padding
      padding = get_pad_tuple2d(padding)
      return _make.contrib_conv2d_winograd_without_weight_transform(
          data,
          weight,
          tile_size,
          strides,
          padding,
          dilation,
          groups,
          channels,
          kernel_size,
          data_layout,
          kernel_layout,
          out_layout,
          out_dtype,
      )
  
  
  def contrib_conv2d_gemm_without_weight_transform(
      data,
      weight,
      strides=(1, 1),
      padding=(0, 0),
      dilation=(1, 1),
      groups=1,
      channels=None,
      kernel_size=None,
      data_layout="NCHW",
      kernel_layout="OIHW",
      out_layout="",
      out_dtype="",
  ):
      r"""2D convolution with gemm algorithm.
  
      The basic parameters are the same as the ones in vanilla conv2d.
      It assumes the weight is pre-transformed by nn.contrib_conv2d_gemm_weight_transform
  
      Parameters
      ----------
      data : tvm.relay.Expr
          The input data to the operator.
  
      weight : tvm.relay.Expr
          The weight expressions.
  
      strides : tuple of int, optional
          The strides of convolution.
  
      padding : tuple of int, optional
          The padding of convolution on both sides of inputs before convolution.
  
      dilation : tuple of int, optional
          Specifies the dilation rate to be used for dilated convolution.
  
      groups : int, optional
          Number of groups for grouped convolution.
  
      channels : int, optional
          Number of output channels of this convolution.
  
      kernel_size : tuple of int, optional
          The spatial of the convolution kernel.
  
      data_layout : str, optional
          Layout of the input.
  
      kernel_layout : str, optional
          Layout of the weight.
  
      out_layout : str, optional
          Layout of the output, by default, out_layout is the same as data_layout
  
      out_dtype : str, optional
          Specifies the output data type for mixed precision conv2d.
  
      Returns
      -------
      result : tvm.relay.Expr
          The computed result.
      """
      # convert 2-way padding to 4-way padding
      padding = get_pad_tuple2d(padding)
      return _make.contrib_conv2d_gemm_without_weight_transform(
          data,
          weight,
          strides,
          padding,
          dilation,
          groups,
          channels,
          kernel_size,
          data_layout,
          kernel_layout,
          out_layout,
          out_dtype,
      )
  
  
  def contrib_conv2d_nchwc(
      data,
      kernel,
      strides=(1, 1),
      padding=(0, 0),
      dilation=(1, 1),
      groups=1,
      channels=None,
      kernel_size=None,
      data_layout="NCHW8c",
      kernel_layout="OIHW",
      out_layout="",
      out_dtype="",
  ):
      r"""Variant of 2D convolution.
  
      This operator takes the weight as the convolution kernel
      and convolves it with data to produce an output, following a specialized
      NCHWc data layout.
  
      Parameters
      ----------
      data : tvm.relay.Expr
          The input data to the operator.
  
      kernel : tvm.relay.Expr
          The kernel expressions.
  
      strides : tuple of int, optional
          The strides of convolution.
  
      padding : tuple of int, optional
          The padding of convolution on both sides of inputs before convolution.
  
      dilation : tuple of int, optional
          Specifies the dilation rate to be used for dilated convolution.
  
      groups : int, optional
          Number of groups for grouped convolution.
  
      channels : int, optional
          Number of output channels of this convolution.
  
      kernel_size : tuple of int, optional
          The spatial of the convolution kernel.
  
      data_layout : str, optional
          Layout of the input.
  
      kernel_layout : str, optional
          Layout of the weight.
  
      out_layout : str, optional
          Layout of the output, by default, out_layout is the same as data_layout
  
      out_dtype : str, optional
          Specifies the output data type for mixed precision conv2d.
  
      Returns
      -------
      result : tvm.relay.Expr
          The computed result.
      """
      # convert 2-way padding to 4-way padding
      padding = get_pad_tuple2d(padding)
      return _make.contrib_conv2d_NCHWc(
          data,
          kernel,
          strides,
          padding,
          dilation,
          groups,
          channels,
          kernel_size,
          data_layout,
          kernel_layout,
          out_layout,
          out_dtype,
      )
  
  
  def contrib_depthwise_conv2d_nchwc(
      data,
      kernel,
      strides=(1, 1),
      padding=(0, 0),
      dilation=(1, 1),
      groups=1,
      channels=None,
      kernel_size=None,
      data_layout="NCHW8c",
      kernel_layout="OIHW",
      out_layout="",
      out_dtype="",
  ):
      r"""Variant of 2D depthwise convolution.
  
      This operator takes the weight as the depthwise convolution kernel
      and depthwise convolves it with data to produce an output, following a specialized
      NCHWc data layout.
  
      Parameters
      ----------
      data : tvm.relay.Expr
          The input data to the operator.
  
      kernel : tvm.relay.Expr
          The kernel expressions.
  
      strides : tuple of int, optional
          The strides of convolution.
  
      padding : tuple of int, optional
          The padding of convolution on both sides of inputs before convolution.
  
      dilation : tuple of int, optional
          Specifies the dilation rate to be used for dilated convolution.
  
      groups : int, optional
          Number of groups for grouped convolution.
  
      channels : int, optional
          Number of output channels of this convolution.
  
      kernel_size : tuple of int, optional
          The spatial of the convolution kernel.
  
      data_layout : str, optional
          Layout of the input.
  
      kernel_layout : str, optional
          Layout of the weight.
  
      out_layout : str, optional
          Layout of the output, by default, out_layout is the same as data_layout
  
      out_dtype : str, optional
          Specifies the output data type for mixed precision conv2d.
  
      Returns
      -------
      result : tvm.relay.Expr
          The computed result.
      """
      # convert 2-way padding to 4-way padding
      padding = get_pad_tuple2d(padding)
      return _make.contrib_depthwise_conv2d_NCHWc(
          data,
          kernel,
          strides,
          padding,
          dilation,
          groups,
          channels,
          kernel_size,
          data_layout,
          kernel_layout,
          out_layout,
          out_dtype,
      )
  
  
  def contrib_conv2d_winograd_weight_transform(weight, tile_size):
      r"""Weight Transformation part for 2D convolution with winograd algorithm.
  
      We separate this as a single op to enable pre-compute for inference.
      Use this together with nn.contrib_conv2d_winograd_without_weight_transform
  
      Parameters
      ----------
      weight : tvm.relay.Expr
          The weight expressions.
  
      tile_size : int
          The Tile size of winograd. E.g. 2 for F(2x2, 3x3) and 4 for F(4x4, 3x3)
  
      Returns
      -------
      result : tvm.relay.Expr
          The computed result.
      """
      return _make.contrib_conv2d_winograd_weight_transform(weight, tile_size)
  
  
  def contrib_conv2d_gemm_weight_transform(weights, tile_rows, tile_cols):
      r"""Weight Transformation part for 2D convolution with gemm algorithm.
  
      We separate this as a single op to enable pre-compute for inference.
      Use this together with nn.contrib_conv2d_gemm_without_weight_transform
  
      Parameters
      ----------
      weights : tvm.relay.Expr
          The weight expressions.
      tile_rows: int
          Tile rows of the weight transformation for ConvGemm.
      tile_cols: int
         Tile columns of the weight transformation for ConvGemm.
  
      Returns
      -------
      result : tvm.relay.Expr
          The computed result.
      """
      return _make.contrib_conv2d_gemm_weight_transform(weights, tile_rows, tile_cols)
  
  
  def contrib_conv3d_winograd_weight_transform(weight, tile_size):
      r"""Weight Transformation part for 3D convolution with winograd algorithm.
  
      We separate this as a single op to enable pre-compute for inference.
      Use this together with nn.contrib_conv3d_winograd_without_weight_transform
  
      Parameters
      ----------
      weight : tvm.relay.Expr
          The weight expressions.
  
      tile_size : int
          The Tile size of winograd. E.g. 2 for F(2x2x2, 3x3x3) and 4 for F(4x4x4, 3x3x3)
  
      Returns
      -------
      result : tvm.relay.Expr
          The computed result.
      """
      return _make.contrib_conv3d_winograd_weight_transform(weight, tile_size)
  
  
  def contrib_conv2d_winograd_nnpack_weight_transform(weight, convolution_algorithm, out_dtype=""):
      r"""Weight Transformation part for 2D convolution with winograd algorithm.
  
      We separate this as a single op to enable pre-compute for inference.
      Use this together with nn.contrib_conv2d_winograd_without_weight_transform
  
      Parameters
      ----------
      weight : tvm.relay.Expr
          The weight expressions.
  
      convolution_algorithm : int
          The Tile size of winograd. E.g. 2 for F(2x2, 3x3) and 4 for F(4x4, 3x3)
  
      Returns
      -------
      result : tvm.relay.Expr
          The computed result.
      """
      return _make.contrib_conv2d_winograd_nnpack_weight_transform(
          weight, convolution_algorithm, out_dtype
      )
  
  
  def deformable_conv2d(
      data,
      offset,
      weight,
      strides=(1, 1),
      padding=(0, 0),
      dilation=(1, 1),
      deformable_groups=1,
      groups=1,
      channels=None,
      kernel_size=None,
      data_layout="NCHW",
      kernel_layout="OIHW",
      out_layout="",
      out_dtype="",
  ):
      r"""Deformable 2d convolution.
  
      The deformable convolution operation is described in https://arxiv.org/abs/1703.06211
  
      Parameters
      ----------
      data : tvm.relay.Expr
          The input data to the operator.
  
      offset : tvm.relay.Expr
          The offset expressions.
  
      weight : tvm.relay.Expr
          The weight expressions.
  
      strides : tuple of int, optional
          The strides of convolution.
  
      padding : tuple of int, optional
          The padding of convolution on both sides of inputs before convolution.
  
      dilation : tuple of int, optional
          Specifies the dilation rate to be used for dilated convolution.
  
      deformable_groups : int, optional
          Number of deformable groups.
  
      groups : int, optional
          Number of groups for grouped convolution.
  
      channels : int, optional
          Number of output channels of this convolution.
  
      kernel_size : tuple of int, optional
          The spatial of the convolution kernel.
  
      data_layout : str, optional
          Layout of the input.
  
      kernel_layout : str, optional
          Layout of the weight.
  
      out_layout : str, optional
          Layout of the output, by default, out_layout is the same as data_layout
  
      out_dtype : str, optional
          Specifies the output data type for mixed precision conv2d.
  
      Returns
      -------
      result : tvm.relay.Expr
          The computed result.
  
      """
      # convert 2-way padding to 4-way padding
      padding = get_pad_tuple2d(padding)
      return _make.deformable_conv2d(
          data,
          offset,
          weight,
          strides,
          padding,
          dilation,
          deformable_groups,
          groups,
          channels,
          kernel_size,
          data_layout,
          kernel_layout,
          out_layout,
          out_dtype,
      )
  
  
  def bitpack(data, bits=1, pack_axis=1, bit_axis=2, pack_type="uint32", name="BitPack"):
      """Tensor packing for bitserial operations.
  
      The values along the input tensor's pack_axis are quantized
      and packed together into the specified pack_type in a new bit axis.
  
      For example, consider bitpacking with data to be a tensor with shape `[1, 64, 128, 128]`,
      pack_axis=1, bit_axis=4, pack_type=uint8, and bits=2. The output in this case will
      be of shape `[1, 8, 128, 128, 2]`. The dimension of axis 1 has been reduced by a factor
      of 8 since each value is packed into an 8-bit uint8. Axis 4 is now two bitplanes
      representing the quantized value of the incoming data. The output tensor is now
      ready to be used in a bitserial operation.
  
      Parameters
      ----------
      data : tvm.relay.expr
          The incoming tensor to be packed.
  
      bits : int
          Number of bits that should be packed.
  
      pack_axis : int
          Axis that should be decomposed and packed.
  
      bit_axis : int
          New axis containing bitplane.
  
      pack_type : str
          Datatype to pack bits into.
  
      name : str, optional
          Name of the operation.
  
      Returns
      -------
      result : tvm.relay.Expr
          The packed tensor.
      """
      return _make.bitpack(data, bits, pack_axis, bit_axis, pack_type, name)
  
  
  def bitserial_conv2d(
      data,
      weight,
      strides=(1, 1),
      padding=(0, 0),
      channels=None,
      kernel_size=(3, 3),
      activation_bits=1,
      weight_bits=1,
      data_layout="NCHW",
      kernel_layout="OIHW",
      pack_dtype="uint32",
      out_dtype="int16",
      unipolar=True,
  ):
      r"""2D convolution using bitserial computation.
  
      Parameters
      ----------
      data : tvm.relay.Expr
          The input data to the operator.
  
      weight : tvm.relay.Expr
          The weight expressions.
  
      strides : tuple of int, optional
          The strides of convolution.
  
      padding : tuple of int, optional
          The padding of convolution on both sides of inputs before convolution.
  
      channels : int, optional
          Number of output channels of this convolution.
  
      kernel_size : tuple of int, optional
          The spatial of the convolution kernel.
  
      activation_bits : int
          Number of bits to pack for activations.
  
      weight_bits : int
          Number of bits to pack for weights.
  
      data_layout : str, optional
          Layout of the input.
  
      kernel_layout : str, optional
          Layout of the kernel
  
      pack_dtype: str, optional
          Datatype to pack bits into.
  
      out_dtype : str, optional
          Specifies the output data type for mixed precision conv2d.
  
      Returns
      -------
      result : tvm.relay.Expr
          The computed result.
      """
      # convert 2-way padding to 4-way padding
      padding = get_pad_tuple2d(padding)
      return _make.bitserial_conv2d(
          data,
          weight,
          strides,
          padding,
          channels,
          kernel_size,
          activation_bits,
          weight_bits,
          data_layout,
          kernel_layout,
          pack_dtype,
          out_dtype,
          unipolar,
      )
  
  
  def bitserial_dense(
      data,
      weight,
      units=None,
      data_bits=1,
      weight_bits=1,
      pack_dtype="uint32",
      out_dtype="int16",
      unipolar=True,
  ):
      """Bitserial Dense operator.
      Applies matrix multiplication of two quantized matrices
      using a fast bitserial algorithm.
  
      .. math::
  
      `Y = X * W`
  
      Parameters
      ----------
      data : tvm.relay.Expr
          The input data to the operator.
  
      weight : tvm.relay.Expr
          The weight expressions.
  
      units : int, optional
          Number of hidden units of the dense transformation.
  
      data_bits : int
          Number of bits incoming tensor should be packed with.
  
      weight_bits : int
          Number of bits weight tensor should be packed with.
  
      pack_dtype : str, optional
          Datatype to pack individual bits into before computation.
  
      out_dtype : str, optional
          Specifies the output data type for mixed precision dense.
  
      unipolar : bool, optional
          Whether to use unipolar or bipolar quantization for inputs.
  
      Returns
      -------
      result : tvm.relay.Expr
          The computed result.
      """
      return _make.bitserial_dense(
          data, weight, units, data_bits, weight_bits, pack_dtype, out_dtype, unipolar
      )
  
  
  def cross_entropy(predictions, targets):
      """CrossEntropy without logits.
  
      Parameters
      ----------
      predictions : tvm.relay.Expr
        The predictions.
  
      targets : tvm.relay.Expr
        The targets.
  
      Returns
      -------
      result : tvm.relay.Expr
        The computed result.
      """
      return _make.cross_entropy(predictions, targets)
  
  
  def cross_entropy_with_logits(predictions, targets):
      """CrossEntropy with logits.
  
      Parameters
      ----------
      predictions : tvm.relay.Expr
        The predictions.
  
      targets : tvm.relay.Expr
        The targets.
  
      Returns
      -------
      result : tvm.relay.Expr
        The computed result.
      """
      return _make.cross_entropy_with_logits(predictions, targets)
  
  
  def depth_to_space(data, block_size, layout="NCHW", mode="DCR"):
      """Convert channels into spatial blocks.
  
      Parameters
      ----------
      data : tvm.relay.Expr
          Input data with channels divisible by block_size**2
  
      block_size : int
          Size of blocks to convert channels into.
  
      layout : string
          One of NCHW or NHWC, indicates channel axis.
  
      mode : string
          One of DCR or CDR, indicates which order channels
          are accessed in.
  
      Returns
      -------
      result : tvm.relay.Expr
          Tensor with shape [in_batch, in_channel / block_size * block_size,
                             in_height * block_size, in_width * block_size]
      """
      return _make.depth_to_space(data, block_size, layout, mode)
  
  
  def space_to_depth(data, block_size, layout="NCHW"):
      """Convert spatial blocks into channels.
  
      Parameters
      ----------
      data : tvm.relay.Expr
          Input data with spatial dimensions divisible by block_size
  
      block_size : int
          Size of blocks to decompose into channels.
  
      layout : string
          One of NCHW or NHWC, indicates channel axis.
  
      Returns
      -------
      result : tvm.relay.Expr
          Tensor with shape [in_batch, in_channel * block_size * block_size,
                             in_height / block_size, in_width / block_size]
      """
      return _make.space_to_depth(data, block_size, layout)
  
  
  def adaptive_max_pool2d(data, output_size=None, layout="NCHW"):
      r"""2D adaptive max pooling operator. This operator is experimental.
  
      This operator takes data as input and does 2D max 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 : tvm.relay.Expr
          The input data to the operator.
  
      output_size : tuple of int. optional
          Output height and width.
  
      layout : str, optional
          Layout of the input.
  
      Returns
      -------
      result : tvm.relay.Expr
          The computed result.
      """
      output_size = [] or output_size
      return _make.adaptive_max_pool2d(data, output_size, layout)
  
  
  def adaptive_avg_pool2d(data, output_size=None, layout="NCHW"):
      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 : tvm.relay.Expr
          The input data to the operator.
  
      output_size : tuple of int. optional
          Output height and width.
  
      layout : str, optional
          Layout of the input.
  
      Returns
      -------
      result : tvm.relay.Expr
          The computed result.
      """
      output_size = [] or output_size
      return _make.adaptive_avg_pool2d(data, output_size, layout)
  
  
  def adaptive_max_pool3d(data, output_size=None, layout="NCDHW"):
      r"""3D adaptive max pooling operator. This operator is experimental.
  
      This operator takes data as input and does 3D max value calculation
      across each window represented by DxWxH.
  
      In the default case, where the data_layout is `NCDHW`
      a data Tensor with shape `(batch_size, in_channels, depth, height, width)`,
      to produce an output Tensor with shape
      (batch_size, in_channels, output_depth, 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 depth, height and width will be used
          as output depth, 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 x output_size) for any input (NCDHW).
  
          If a tuple of integers (depth, height, width) are provided for output_size,
          the output size is (N x C x depth x height x width) for any input (NCDHW).
  
      Parameters
      ----------
      data : tvm.relay.Expr
          The input data to the operator.
  
      output_size : tuple of int. optional
          Output height and width.
  
      layout : str, optional
          Layout of the input.
  
      Returns
      -------
      result : tvm.relay.Expr
          The computed result.
      """
      output_size = [] or output_size
      return _make.adaptive_max_pool3d(data, output_size, layout)
  
  
  def adaptive_avg_pool3d(data, output_size=None, layout="NCDHW"):
      r"""3D adaptive avg pooling operator. This operator is experimental.
  
      This operator takes data as input and does 3D avg value calculation
      across each window represented by DxWxH.
  
      In the default case, where the data_layout is `NCDHW`
      a data Tensor with shape `(batch_size, in_channels, depth, height, width)`,
      to produce an output Tensor with shape
      (batch_size, in_channels, output_depth, 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 depth, height and width will be used
          as output depth, 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 x output_size) for any input (NCDHW).
  
          If a tuple of integers (depth, height, width) are provided for output_size,
          the output size is (N x C x depth x height x width) for any input (NCDHW).
  
      Parameters
      ----------
      data : tvm.relay.Expr
          The input data to the operator.
  
      output_size : tuple of int. optional
          Output height and width.
  
      layout : str, optional
          Layout of the input.
  
      Returns
      -------
      result : tvm.relay.Expr
          The computed result.
      """
      output_size = [] or output_size
      return _make.adaptive_avg_pool3d(data, output_size, layout)
  
  
  def global_max_pool3d(data, layout="NCDHW"):
      r"""3D global maximum pooling operator.
  
      This operator takes data as input and does 3D max value calculation
      across each window represented by DxWxH.
  
      In the default case, where the data_layout is `NCDHW`
      a data Tensor with shape `(batch_size, in_channels, depth, height, width)`,
      to produce an output Tensor with the following rule:
  
      with data of shape (b, c, d, h, w)
      .. math::
  
          \mbox{out}(b, c, 1, 1, 1)  =  \max_{l=0, \ldots, d},  \max_{m=0, \ldots, h},
               \max_{n=0, \ldots, w} \mbox{data}(b, c, l, m, n)
  
      Parameters
      ----------
      data : tvm.relay.Expr
          The input data to the operator.
  
      layout : str, optional
          Layout of the input.
  
      Returns
      -------
      result : tvm.relay.Expr
          The computed result.
      """
      output_size = [1, 1, 1]
      return _make.adaptive_max_pool3d(data, output_size, layout)
  
  
  def global_avg_pool3d(data, layout="NCDHW"):
      r"""3D global average pooling operator.
  
      This operator takes data as input and does 3D average value calculation
      across each window represented by DxWxH.
  
      In the default case, where the data_layout is `NCDHW`
      a data Tensor with shape `(batch_size, in_channels, depth, height, width)`,
      to produce an output Tensor with the following rule:
  
      with data of shape (b, c, d, h, w)
  
      .. math::
  
          \mbox{out}(b, c, 1, 1, 1)  = \frac{1}{d * h * w} \sum_{l=0}^{d-1}  \sum_{m=0}^{h-1}
               \sum_{n=0}^{w-1} \mbox{data}(b, c, l, m, n)
  
      Parameters
      ----------
      data : tvm.relay.Expr
          The input data to the operator.
  
      layout : str, optional
          Layout of the input.
  
      Returns
      -------
      result : tvm.relay.Expr
          The computed result.
      """
      output_size = [1, 1, 1]
      return _make.adaptive_avg_pool3d(data, output_size, layout)
  
  
  def correlation(
      data1, data2, kernel_size, max_displacement, stride1, stride2, padding, is_multiply, layout
  ):
      r"""Applies correlation to inputs.
  
      The correlation layer performs multiplicative patch comparisons between two feature maps.
      Given two multi-channel feature maps :math:`f_{1}, f_{2}`, with :math:`w`, :math:`h`, and
      :math:`c` being their width, height, and number of channels, the correlation layer lets the
      network compare each patch from :math:`f_{1}` with each patch from :math:`f_{2}`.
  
      For now we consider only a single comparison of two patches. The 'correlation' of two patches
      centered at :math:`x_{1}` in the first map and :math:`x_{2}` in the second map is then defined
      as:
  
      .. math::
  
          c(x_{1}, x_{2}) = \sum_{o \in [-k,k] \times [-k,k]} <f_{1}(x_{1} + o), f_{2}(x_{2} + o)>
  
      for a square patch of size :math:`K:=2k+1`.
  
      Note that the equation above is identical to one step of a convolution in neural networks, but
      instead of convolving data with a filter, it convolves data with other    data. For this
      reason, it has no training weights.
  
      Computing :math:`c(x_{1}, x_{2})` involves :math:`c * K^{2}` multiplications. Comparing all
      patch combinations involves :math:`w^{2}*h^{2}` such computations.
  
      Given a maximum displacement :math:`d`, for each location :math:`x_{1}` it computes
      correlations :math:`c(x_{1}, x_{2})` only in a neighborhood of size :math:`D:=2d+1`,
      by limiting the range of :math:`x_{2}`. We use strides :math:`s_{1}, s_{2}`, to quantize
      :math:`x_{1}` globally and to quantize :math:`x_{2}` within the neighborhood
      centered around :math:`x_{1}`.
  
      The final output is defined by the following expression:
  
      .. math::
  
          out[n, q, i, j] = c(x_{i, j}, x_{q})
  
      where :math:`i` and :math:`j` enumerate spatial locations in :math:`f_{1}`, and :math:`q`
      denotes the :math:`q^{th}` neighborhood of :math:`x_{i,j}`.
  
      Parameters
      ----------
      data1 : tvm.te.Tensor
          4-D with shape [batch, channel, height, width]
  
      data2 : tvm.te.Tensor
          4-D with shape [batch, channel, height, width]
  
      kernel_size: int
          Kernel size for correlation, must be an odd number
  
      max_displacement: int
          Max displacement of Correlation
  
      stride1: int
          Stride for data1
  
      stride2: int
          Stride for data2 within the neightborhood centered around data1
  
      padding : int or a list/tuple of 2 or 4 ints
          Padding size, or
          [pad_height, pad_width] for 2 ints, or
          [pad_top, pad_left, pad_bottom, pad_right] for 4 ints
  
      is_multiply: bool
          operation type is either multiplication or substraction
  
      layout: str
          layout of data1, data2 and the output
  
      Returns
      -------
      Output : tvm.te.Tensor
          4-D with shape [batch, out_channel, out_height, out_width]
      """
      if isinstance(padding, int):
          padding = (padding, padding)
      return _make.correlation(
          data1, data2, kernel_size, max_displacement, stride1, stride2, padding, is_multiply, layout
      )
  
  
  def space_to_batch_nd(data, block_shape, paddings, pad_value=0):
      r"""Divide spatial dimensions of the data into a grid of blocks
      and interleave them into batch dim.
  
      Parameters
      ----------
      data : tvm.te.Tensor
          N-D with shape [batch, spatial_shape, remaining_shape]
  
      block_shape : relay.Expr
          1-D of size [M] where M is number of spatial dims, specifies block size
          for each spatial dimension.
  
      paddings : relay.Expr
          2-D of shape [M, 2] where M is number of spatial dims, specifies
          [before, after] paddings for each spatial dimension.
  
      pad_value : float, or relay.Expr, optional, default=0
          The value used for padding.
  
      Returns
      -------
      result : relay.Expr
          N-D Tensor with shape
          [in_batch * prod(block_shape),
          padded_data[1] / block_shape[0], ..., padded_data[M] / block_shape[M-1],
          remaining_shape]
      """
  
      return _make.space_to_batch_nd(data, block_shape, paddings, pad_value)
  
  
  def batch_to_space_nd(data, block_shape, crops):
      r"""Reshape the batch dimension into spatial dimensions.
  
      Parameters
      ----------
      data : tvm.te.Tensor
          N-D with shape [batch, spatial_shape, remaining_shape]
  
      block_shape : relay.Expr
          1-D of size [M] where M is number of spatial dims, specifies block size
          for each spatial dimension.
  
      crops : relay.Expr
          2-D of shape [M, 2] where M is number of spatial dims, specifies
          [begin, end] crop size for each spatial dimension.
  
      Returns
      -------
      result : relay.Expr
          N-D Tensor with shape
          [batch / prod(block_shape),
          in_shape[1] * block_shape[0] - crops[0,0] - crops[0,1], ...,
          in_shape[M] * block_shape[M-1] - crops[M-1, 0] - crops[M-1, 1],
          remaining_shape]
      """
  
      return _make.batch_to_space_nd(data, block_shape, crops)
  