Add spec: max_pool3d_with_indices

Reviewed By: zonglinpeng

Differential Revision: D76528169

fbshipit-source-id: 0c7602f6c0d75930e1b7810bb4aa8fd8a3e9d054

Author: manuelcandales

facto/specdb/__init__.py

@@ -2500,6 +2500,69 @@
             OutArg(ArgType.Tensor, name="indices"),
         ],
     ),
+    Spec(  # TODO(mcandales): Calibrate.
+        op="max_pool3d_with_indices.default",  # (Tensor self, int[3] kernel_size, int[3] stride=[], int[3] padding=0, int[3] dilation=1, bool ceil_mode=False) -> (Tensor, Tensor)
+        inspec=[
+            InPosArg(  # self
+                ArgType.Tensor,
+                name="self",
+                deps=[1, 2, 3, 4, 5],
+                constraints=[
+                    cp.Dtype.Ne(lambda deps: torch.bool),
+                    cp.Rank.In(lambda deps: [4, 5]),
+                    # cp.Size.Ge(lambda deps, r, d: 0 if d == 0 and r == 5 else 1),
+                    cp.Size.Ge(
+                        lambda deps, r, d: fn.pool_input_size_min(
+                            3, deps[0], deps[1], deps[2], deps[3], deps[4], r, d
+                        )
+                    ),
+                ],
+            ),
+            InPosArg(  # kernel_size
+                ArgType.LengthList,
+                name="kernel_size",
+                constraints=[
+                    cp.Length.In(lambda deps: [1, 3]),
+                    cp.Value.Ge(lambda deps, length, ix: 1),
+                ],
+            ),
+            InPosArg(  # stride
+                ArgType.LengthList,
+                name="stride",
+                constraints=[
+                    cp.Length.In(lambda deps: [0, 1, 3]),
+                    cp.Value.Ge(lambda deps, length, ix: 1),
+                ],
+            ),
+            InPosArg(  # padding
+                ArgType.LengthList,
+                name="padding",
+                deps=[1],
+                constraints=[
+                    cp.Length.In(lambda deps: [1, 3]),
+                    cp.Value.Ge(lambda deps, length, ix: 0),
+                    cp.Value.Le(
+                        lambda deps, length, ix: fn.pool_padding_max(
+                            deps[0], length, ix
+                        )
+                    ),
+                ],
+            ),
+            InPosArg(  # dilation
+                ArgType.LengthList,
+                name="dilation",
+                constraints=[
+                    cp.Length.In(lambda deps: [1, 3]),
+                    cp.Value.Ge(lambda deps, length, ix: 1),
+                ],
+            ),
+            InPosArg(ArgType.Bool, name="ceil_mode"),  # ceil_mode
+        ],
+        outspec=[
+            OutArg(ArgType.Tensor, name="out"),
+            OutArg(ArgType.Tensor, name="indices"),
+        ],
+    ),
     Spec(
         op="maximum.default",  # (Tensor self, Tensor other) -> Tensor
         inspec=[

facto/specdb/db.py

@@ -459,8 +459,8 @@ def pool_input_size_min(
     kernel_ndim, kernel_size, stride, padding, dilation, ceil_mode, rank, dim
 ):
     if dim == 0:
-        return 0 if rank == 4 else 1
-    if dim == 1 and rank == 4:
+        return 0 if rank == (kernel_ndim + 2) else 1
+    if dim == 1 and rank == (kernel_ndim + 2):
         return 1
 
     kdim = dim - (rank - kernel_ndim)