Move Ops in relay.op.contrib.* (#4942)

* move contrib

* lint

* address comment

* address comment

python/tvm/relay/frontend/mxnet.py

@@ -313,7 +313,7 @@ def _pool2d(new_op, is_avg):
 
 def _mx_adaptive_avg_pooling(inputs, attrs):
     output_size = attrs.get_int_tuple("output_size", [])
-    return _op.contrib.adaptive_avg_pool2d(inputs[0], output_size)
+    return _op.nn.adaptive_avg_pool2d(inputs[0], output_size)
 
 
 def _mx_dropout(inputs, attrs):

python/tvm/relay/frontend/pytorch.py

@@ -151,7 +151,7 @@ def _impl(inputs, input_types):
         data = inputs[0]
         output_size = _infer_shape(inputs[1])
 
-        return _op.contrib.contrib.adaptive_avg_pool2d(
+        return _op.nn.adaptive_avg_pool2d(
             data,
             output_size=output_size)
     return _impl
@@ -161,7 +161,7 @@ def _impl(inputs, input_types):
         data = inputs[0]
         output_size = _infer_shape(inputs[1])
 
-        return _op.contrib.contrib.adaptive_max_pool2d(
+        return _op.nn.adaptive_max_pool2d(
             data,
             output_size=output_size)
     return _impl

python/tvm/relay/op/__init__.py

@@ -32,7 +32,6 @@
 from . import memory
 from . import image
 from . import vision
-from . import contrib
 from . import op_attrs
 
 

python/tvm/relay/op/_tensor.py

@@ -70,6 +70,8 @@
 register_injective_schedule("right_shift")
 register_injective_schedule("left_shift")
 register_injective_schedule("shape_of")
+register_injective_schedule("ndarray_size")
+
 
 # zeros
 @register_compute("zeros")

python/tvm/relay/op/contrib/__init__.py

@@ -17,5 +17,3 @@
 # pylint: disable=wildcard-import
 """Neural network related operators."""
 from __future__ import absolute_import as _abs
-from .contrib import *
-from . import _contrib

python/tvm/relay/op/contrib/contrib.py

@@ -17,115 +17,3 @@
 #pylint: disable=invalid-name, too-many-lines
 """Contrib operations."""
 from __future__ import absolute_import as _abs
-from . import _make
-
-
-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 ndarray_size(data, dtype="int32"):
-    """Get number of elements of input tensor.
-
-    Parameters
-    ----------
-    data : tvm.relay.Expr
-        The input tensor.
-
-    dtype : str, optional
-        The target data type.
-
-    Returns
-    -------
-    result : tvm.relay.Expr
-        The number of elements of input tensor.
-    """
-    return _make.ndarray_size(data, dtype)

python/tvm/relay/op/nn/_nn.py

@@ -247,6 +247,16 @@ def legalize_conv2d_transpose(attrs, inputs, types):
 reg.register_pattern("nn.global_avg_pool2d", OpPattern.OUT_ELEMWISE_FUSABLE)
 
 
+# adaptive_max_pool2d
+reg.register_schedule("nn.adaptive_max_pool2d", strategy.schedule_adaptive_pool)
+reg.register_pattern("nn.adaptive_max_pool2d", OpPattern.OUT_ELEMWISE_FUSABLE)
+
+
+# adaptive_avg_pool2d
+reg.register_schedule("nn.adaptive_avg_pool2d", strategy.schedule_adaptive_pool)
+reg.register_pattern("nn.adaptive_avg_pool2d", OpPattern.OUT_ELEMWISE_FUSABLE)
+
+
 # leaky_relu
 reg.register_broadcast_schedule("nn.leaky_relu")
 reg.register_pattern("nn.leaky_relu", OpPattern.ELEMWISE)

python/tvm/relay/op/nn/nn.py

@@ -2277,3 +2277,97 @@ def space_to_depth(data, block_size, layout='NCHW'):
                            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)

python/tvm/relay/op/tensor.py

@@ -974,3 +974,22 @@ def shape_of(data, dtype="int32"):
         The shape tensor.
     """
     return _make.shape_of(data, dtype)
+
+
+def ndarray_size(data, dtype="int32"):
+    """Get number of elements of input tensor.
+
+    Parameters
+    ----------
+    data : tvm.relay.Expr
+        The input tensor.
+
+    dtype : str, optional
+        The target data type.
+
+    Returns
+    -------
+    result : tvm.relay.Expr
+        The number of elements of input tensor.
+    """
+    return _make.ndarray_size(data, dtype)