update doc for 3d

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

@@ -2376,29 +2376,28 @@ def adaptive_avg_pool2d(data,
 def adaptive_max_pool3d(data,
                         output_size=None,
                         layout="NCDHW"):
-    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.
+    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 `NCHW`
-    a data Tensor with shape `(batch_size, in_channels, height, width)`,
+    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_height, output_width).
+    (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 height and width will be used
-        as output height and width.
+        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) for any input (NCHW).
+        (N x C x output_size, output_size x output_size) for any input (NCDHW).
 
-        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).
+        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
     ----------
@@ -2423,29 +2422,28 @@ def adaptive_max_pool3d(data,
 def adaptive_avg_pool3d(data,
                         output_size=None,
                         layout="NCDHW"):
-    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.
+    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 `NCHW`
-    a data Tensor with shape `(batch_size, in_channels, height, width)`,
+    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_height, output_width).
+    (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 height and width will be used
-        as output height and width.
+        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) for any input (NCHW).
+        (N x C x output_size, output_size x output_size) for any input (NCDHW).
 
-        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).
+        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
     ----------

src/relay/op/nn/pooling.cc

@@ -692,16 +692,16 @@ TVM_REGISTER_GLOBAL("relay.op.nn._make.adaptive_max_pool3d")
 RELAY_REGISTER_OP("nn.adaptive_max_pool3d")
   .describe(R"code(Adaptive max pooling operation for 3D data.
 
-- **data**: This depends on the `layout` parameter. Input is 4D array of shape
-            (batch_size, channels, height, width) if `layout` is `NCHW`.
-- **output_size**: If this argument is not provided, input height and width will be used
-                   as output height and width.
+- **data**: This depends on the `layout` parameter. Input is 5D array of shape
+            (batch_size, channels, depth, height, width) if `layout` is `NCDHW`.
+- **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) 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).
-- **out**: This depends on the `layout` parameter. Output is 4D array of shape
-           (batch_size, channels, output_height, output_width)  if `layout` is `NCHW`.
+                   (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, height x width) for any input (NCDHW).
+- **out**: This depends on the `layout` parameter. Output is 5D array of shape
+           (batch_size, channels, output_depth, output_height, output_width)  if `layout` is `NCDHW`.
 
 )code" TVM_ADD_FILELINE)
 .set_attrs_type<AdaptivePool3DAttrs>()
@@ -729,18 +729,16 @@ TVM_REGISTER_GLOBAL("relay.op.nn._make.adaptive_avg_pool3d")
 
 RELAY_REGISTER_OP("nn.adaptive_avg_pool3d")
   .describe(R"code(Adaptive avg pooling operation for 3D data.
-
-- **data**: This depends on the `layout` parameter. Input is 4D array of shape
-            (batch_size, channels, height, width) if `layout` is `NCHW`.
-- **output_size**: If this argument is not provided, input height and width will be used
-                   as output height and width.
+- **data**: This depends on the `layout` parameter. Input is 5D array of shape
+            (batch_size, channels, depth, height, width) if `layout` is `NCDHW`.
+- **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) 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).
-- **out**: This depends on the `layout` parameter. Output is 4D array of shape
-           (batch_size, channels, output_height, output_width)  if `layout` is `NCHW`.
-
+                   (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, height x width) for any input (NCDHW).
+- **out**: This depends on the `layout` parameter. Output is 5D array of shape
+           (batch_size, channels, output_depth, output_height, output_width)  if `layout` is `NCDHW`.
 )code" TVM_ADD_FILELINE)
 .set_attrs_type<AdaptivePool3DAttrs>()
 .set_num_inputs(1)