[Operator] Support advanced options for pooling operators

python/hidet/graph/frontend/torch/register_functions.py

@@ -328,27 +328,23 @@ def unsqueeze(x: Tensor, dim: int):
 def avg_pool2d(
     x: Tensor, kernel_size, stride=None, padding=0, ceil_mode=False, count_include_pad=True, divisor_override=None
 ):
-    if ceil_mode:
-        raise NotImplementedError("ceil_mode=True")
-    if not count_include_pad:
-        raise NotImplementedError("count_include_pad=False")
-    if divisor_override is not None:
-        raise NotImplementedError("divisor_override is not None")
     if stride is None:
         stride = kernel_size
-    y = ops.avg_pool2d(x, kernel_size, stride, padding)
+    y = ops.avg_pool2d(
+        x,
+        kernel_size,
+        stride,
+        padding,
+        ceil_mode=ceil_mode,
+        count_include_pad=count_include_pad,
+        divisor_override=divisor_override,
+    )
     return y
 
 
 @register_function(torch.nn.functional.avg_pool3d)
 def avg_pool3d(x: Tensor, kernel_size, stride, padding, ceil_mode=False, count_include_pad=True, divisor_override=None):
-    if ceil_mode:
-        raise NotImplementedError("ceil_mode=True")
-    if not count_include_pad:
-        raise NotImplementedError("count_include_pad=False")
-    if divisor_override is not None:
-        raise NotImplementedError("divisor_override is not None")
-    y = ops.avg_pool3d(x, kernel_size, stride, padding)
+    y = ops.avg_pool3d(x, kernel_size, stride, padding, ceil_mode, count_include_pad, divisor_override)
     return y
 
 
@@ -1238,7 +1234,7 @@ def isinf(x: Tensor) -> Tensor:
 
 
 @register_function(torch.nn.functional.pad)
-def torch_pad(x: Tensor, pad: Union[Tuple[int], List[int]], mode: str = 'constant', value=0):
+def torch_pad(x: Tensor, pad: Union[Tuple[int, ...], List[int]], mode: str = 'constant', value=0):
     if isinstance(pad, tuple):
         pad = list(pad)
     # Torch's pad list has form [p2left, p2right, p1left, p1right, p0left, p0right]

python/hidet/graph/frontend/torch/register_modules.py

@@ -162,6 +162,13 @@ def __call__(self, x: Tensor) -> Tensor:
         )
 
 
+@register_module(torch.nn.ZeroPad2d)
+class HidetZeroPad2d(HidetModule):
+    def __call__(self, x: Tensor) -> Tensor:
+        assert isinstance(self.mod, torch.nn.ZeroPad2d)
+        return regs.torch_pad(x=x, pad=self.mod.padding, mode='constant', value=0.0)
+
+
 @register_module(torch.nn.Linear)
 class HidetLinear(HidetModule):
     def __init__(self, torch_module: torch.nn.Module):

python/hidet/graph/ops/conv2d/conv2d_gemm.py

@@ -352,11 +352,11 @@ def __init__(
             f"expected groups * WC == C, got groups: {groups}, WC: {WC}, C: {C}; make sure the image is channels last!",
         )
         self._assert(
-            DILX > 0 and DILY > 0 and STRX > 0 and STRY > 0,
+            logical_and(DILX > 0, DILY > 0, STRX > 0, STRY > 0),
             f"dilations and strides must be larger than 0, got strides={(STRY, STRX)}, dilations={(DILY, DILX)}",
         )
         self._assert(parallel_k_parts > 0, "expected parallel_k_parts to be greater than 0")
-        self._assert(H >= KY and W >= KX, "expected image dimensions to be greater than filter dimensions")
+        self._assert(logical_and(H >= KY, W >= KX), "expected image dimensions to be greater than filter dimensions")
 
         OUT_H = (H - DILY * (KY - 1) - 1) // STRY + 1
         OUT_W = (W - DILX * (KX - 1) - 1) // STRX + 1