fix(avgpool2d): preserve pooling semantics end to end

act/back_end/dual_tf/tf_cnn.py

@@ -16,6 +16,7 @@
 import torch.nn.functional as F
 from typing import Tuple, Optional, Dict, Any, List, cast
 from act.back_end.core import Bounds, Layer
+from act.back_end.utils import avgpool2d_denominators, avgpool2d_output_hw, pair_2d
 
 from .tf_forward import (
     LinearBound, Frame,
@@ -304,21 +305,20 @@ def forward_avgpool2d(
 def backward_avgpool2d(L, nu, bounds_dict, preds, M: int = 1, alpha=None):
     """AvgPool2D backward — exact linear transpose.
 
-    AvgPool is mathematically EXACT linear: ``y = (1/k²) · conv(x, ones_k)``,
-    so the backward is mathematically exact (no relaxation):
-    ``nu_in = (1/k²) · conv_transpose(nu_out, ones_k)``, per channel.
+    AvgPool is linear. Scale each output coefficient by its window divisor,
+    then apply the transpose of the unit-kernel sum, per channel.
     contrib = 0 (no bias, no nonlinear gap).
     """
     kernel_size = L.params.get("kernel_size", 2)
     stride = L.params.get("stride", kernel_size)
     padding = L.params.get("padding", 0)
+    ceil_mode = bool(L.params.get("ceil_mode", False))
+    count_include_pad = bool(L.params.get("count_include_pad", True))
+    divisor_override = L.params.get("divisor_override")
     input_shape = L.params.get("input_shape")
-    if isinstance(kernel_size, (list, tuple)): kernel_size = int(kernel_size[0])
-    if isinstance(stride, (list, tuple)): stride = int(stride[0])
-    if isinstance(padding, (list, tuple)): padding = int(padding[0])
-    kernel_size = int(kernel_size)
-    stride = int(stride)
-    padding = int(padding)
+    kernel_size = pair_2d(kernel_size)
+    stride = pair_2d(stride)
+    padding = pair_2d(padding)
 
     if input_shape is None:
         raise ValueError(f"backward_avgpool2d: layer {L.id} missing 'input_shape' param")
@@ -329,8 +329,9 @@ def backward_avgpool2d(L, nu, bounds_dict, preds, M: int = 1, alpha=None):
         c, iH, iW = shape[-3:]
     c = int(c); iH = int(iH); iW = int(iW)
 
-    oH = (iH + 2 * padding - kernel_size) // stride + 1
-    oW = (iW + 2 * padding - kernel_size) // stride + 1
+    oH, oW = avgpool2d_output_hw(
+        (iH, iW), kernel_size, stride, padding, ceil_mode
+    )
 
     BM = nu.shape[0]
     v_flat = nu.flatten(start_dim=1)
@@ -345,17 +346,49 @@ def backward_avgpool2d(L, nu, bounds_dict, preds, M: int = 1, alpha=None):
         v_pad[:, :n] = v_flat
         v_4d = v_pad.view(BM, c, oH, oW)
 
-    avg_weight = torch.full((c, 1, kernel_size, kernel_size),
-                            1.0 / (kernel_size * kernel_size),
-                            dtype=nu.dtype, device=nu.device)
-
-    computed_h = (oH - 1) * stride - 2 * padding + kernel_size
-    output_padding = max(0, iH - computed_h)
+    denominators = avgpool2d_denominators(
+        (iH, iW),
+        (oH, oW),
+        kernel_size,
+        stride,
+        padding,
+        ceil_mode=ceil_mode,
+        count_include_pad=count_include_pad,
+        divisor_override=divisor_override,
+        device=nu.device,
+        dtype=nu.dtype,
+    )
+    v_4d = v_4d / denominators.view(1, 1, oH, oW)
+    avg_weight = torch.ones(
+        (c, 1, kernel_size[0], kernel_size[1]),
+        dtype=nu.dtype,
+        device=nu.device,
+    )
 
-    v_out_4d = F.conv_transpose2d(v_4d, avg_weight, None,
-                                  stride=stride, padding=padding,
-                                  output_padding=output_padding,
-                                  groups=c)
+    computed_h = (oH - 1) * stride[0] - 2 * padding[0] + kernel_size[0]
+    computed_w = (oW - 1) * stride[1] - 2 * padding[1] + kernel_size[1]
+    output_padding = (
+        max(0, iH - computed_h),
+        max(0, iW - computed_w),
+    )
+    if output_padding[0] >= stride[0] or output_padding[1] >= stride[1]:
+        raise ValueError("avgpool2d transpose output padding is inconsistent")
+
+    v_out_4d = F.conv_transpose2d(
+        v_4d,
+        avg_weight,
+        None,
+        stride=stride,
+        padding=padding,
+        output_padding=output_padding,
+        groups=c,
+    )
+    v_out_4d = v_out_4d[..., :iH, :iW]
+    if v_out_4d.shape[-2:] != (iH, iW):
+        v_out_4d = F.pad(
+            v_out_4d,
+            (0, iW - v_out_4d.shape[-1], 0, iH - v_out_4d.shape[-2]),
+        )
     v_out = v_out_4d.flatten(start_dim=1)
     contrib = torch.zeros(BM, dtype=nu.dtype, device=nu.device)
 

act/back_end/dual_tf/tf_forward.py

@@ -808,6 +808,9 @@ def _fwd_avgpool2d(layer: Layer, lb: torch.Tensor, ub: torch.Tensor) -> Tuple[to
     kernel_size = layer.params.get("kernel_size", 2)
     stride = layer.params.get("stride", kernel_size)
     padding = layer.params.get("padding", 0)
+    ceil_mode = bool(layer.params.get("ceil_mode", False))
+    count_include_pad = bool(layer.params.get("count_include_pad", True))
+    divisor_override = layer.params.get("divisor_override")
     input_shape = _shape_list(layer.params.get("input_shape"))
     if input_shape is None:
         raise ValueError(
@@ -820,8 +823,16 @@ def _fwd_avgpool2d(layer: Layer, lb: torch.Tensor, ub: torch.Tensor) -> Tuple[to
     else:
         c, h, w = shape[-3], shape[-2], shape[-1]
     B = lb.shape[0]
-    lb_out = F.avg_pool2d(lb.view(B, c, h, w), kernel_size, stride, padding)
-    ub_out = F.avg_pool2d(ub.view(B, c, h, w), kernel_size, stride, padding)
+    pool_kwargs = {
+        "kernel_size": kernel_size,
+        "stride": stride,
+        "padding": padding,
+        "ceil_mode": ceil_mode,
+        "count_include_pad": count_include_pad,
+        "divisor_override": divisor_override,
+    }
+    lb_out = F.avg_pool2d(lb.view(B, c, h, w), **pool_kwargs)
+    ub_out = F.avg_pool2d(ub.view(B, c, h, w), **pool_kwargs)
     return lb_out.flatten(start_dim=1), ub_out.flatten(start_dim=1)
 
 

act/back_end/hybridz_tf/tf_cnn.py

@@ -24,6 +24,7 @@
 from act.back_end.core import Bounds, Fact
 from act.back_end.solver.solver_hz import HZono, SparseHZono, sparse_hz_linear
 from act.back_end.hybridz_tf.tf_mlp import _hz_fact
+from act.back_end.utils import avgpool2d_denominators, avgpool2d_output_hw
 import act.back_end.interval_tf.tf_cnn as interval
 
 
@@ -180,9 +181,23 @@ def sparse_avgpool2d_matrix_from_layer(layer):
     raw_stride = layer.params.get("stride")
     stride = _pair(raw_stride if raw_stride is not None else layer.params["kernel_size"])
     padding = _pair(layer.params.get("padding", 0))
-    OH = (H + 2 * padding[0] - kernel[0]) // stride[0] + 1
-    OW = (W + 2 * padding[1] - kernel[1]) // stride[1] + 1
-    denom = float(kernel[0] * kernel[1])
+    ceil_mode = bool(layer.params.get("ceil_mode", False))
+    count_include_pad = bool(layer.params.get("count_include_pad", True))
+    divisor_override = layer.params.get("divisor_override")
+    OH, OW = avgpool2d_output_hw(
+        (H, W), kernel, stride, padding, ceil_mode
+    )
+    denominators = avgpool2d_denominators(
+        (H, W),
+        (OH, OW),
+        kernel,
+        stride,
+        padding,
+        ceil_mode=ceil_mode,
+        count_include_pad=count_include_pad,
+        divisor_override=divisor_override,
+        dtype=torch.float64,
+    ).cpu().numpy()
     rows, cols, data = [], [], []
     for c in range(C):
         for oh in range(OH):
@@ -198,7 +213,7 @@ def sparse_avgpool2d_matrix_from_layer(layer):
                             continue
                         rows.append(r)
                         cols.append(c * H * W + ih * W + iw)
-                        data.append(1.0 / denom)
+                        data.append(1.0 / float(denominators[oh, ow]))
     return sp.csr_matrix((data, (rows, cols)), shape=(C * OH * OW, C * H * W)), None
 
 
@@ -346,12 +361,25 @@ def _hz_spatial_affine(hz: HZono, op_fn, input_shape, bias=None) -> HZono:
     )
 
 
-def hz_avgpool2d(hz, kernel_size, stride, padding, input_shape) -> HZono:
+def hz_avgpool2d(
+    hz,
+    kernel_size,
+    stride,
+    padding,
+    input_shape,
+    *,
+    ceil_mode=False,
+    count_include_pad=True,
+    divisor_override=None,
+) -> HZono:
     op = lambda x: F.avg_pool2d(
         x,
         kernel_size=kernel_size,
         stride=stride if stride is not None else kernel_size,
         padding=padding,
+        ceil_mode=ceil_mode,
+        count_include_pad=count_include_pad,
+        divisor_override=divisor_override,
     )
     return _hz_spatial_affine(hz, op, input_shape)
 
@@ -367,6 +395,9 @@ def tf_avgpool2d(L, bounds, tf):
                 L.params.get("stride"),
                 L.params.get("padding", 0),
                 ishape,
+                ceil_mode=bool(L.params.get("ceil_mode", False)),
+                count_include_pad=bool(L.params.get("count_include_pad", True)),
+                divisor_override=L.params.get("divisor_override"),
             )
         else:
             hz_in = None

act/back_end/interval_tf/tf_cnn.py

@@ -14,7 +14,7 @@
 import torch.nn.functional as F
 from typing import Callable, Tuple
 from act.back_end.core import Bounds, Con, ConSet, Fact, Layer
-from act.back_end.utils import split_weight
+from act.back_end.utils import avgpool2d_denominators, split_weight
 
 
 def tf_conv2d(L: Layer, Bin: Bounds) -> Fact:
@@ -329,6 +329,9 @@ def tf_avgpool2d(L: Layer, Bin: Bounds) -> Fact:
     kernel_size = L.params["kernel_size"]
     stride = L.params.get("stride", kernel_size)
     padding = L.params.get("padding", 0)
+    ceil_mode = bool(L.params.get("ceil_mode", False))
+    count_include_pad = bool(L.params.get("count_include_pad", True))
+    divisor_override = L.params.get("divisor_override")
 
     # Input/output shape information
     input_shape = L.params["input_shape"]
@@ -342,8 +345,16 @@ def tf_avgpool2d(L: Layer, Bin: Bounds) -> Fact:
 
     input_lb = Bin.lb.view(B_in, channels, in_h, in_w)
     input_ub = Bin.ub.view(B_in, channels, in_h, in_w)
-    output_lb = F.avg_pool2d(input_lb, kernel_size, stride, padding)
-    output_ub = F.avg_pool2d(input_ub, kernel_size, stride, padding)
+    pool_kwargs = {
+        "kernel_size": kernel_size,
+        "stride": stride,
+        "padding": padding,
+        "ceil_mode": ceil_mode,
+        "count_include_pad": count_include_pad,
+        "divisor_override": divisor_override,
+    }
+    output_lb = F.avg_pool2d(input_lb, **pool_kwargs)
+    output_ub = F.avg_pool2d(input_ub, **pool_kwargs)
     assert tuple(output_lb.shape) == (B_in, channels, out_h, out_w), (
         f"avgpool2d output shape mismatch: got {tuple(output_lb.shape)}, expected {(B_in, channels, out_h, out_w)}"
     )
@@ -353,7 +364,16 @@ def tf_avgpool2d(L: Layer, Bin: Bounds) -> Fact:
     B_output = Bounds(output_lb.reshape(B_in, -1), output_ub.reshape(B_in, -1))
 
     W_equiv = _avgpool2d_to_linear_matrix(
-        input_shape, output_shape, kernel_size, stride, padding
+        input_shape,
+        output_shape,
+        kernel_size,
+        stride,
+        padding,
+        ceil_mode=ceil_mode,
+        count_include_pad=count_include_pad,
+        divisor_override=divisor_override,
+        device=Bin.lb.device,
+        dtype=Bin.lb.dtype,
     )
 
     # Create constraints
@@ -364,6 +384,9 @@ def tf_avgpool2d(L: Layer, Bin: Bounds) -> Fact:
         "kernel_size": kernel_size,
         "stride": stride,
         "padding": padding,
+        "ceil_mode": ceil_mode,
+        "count_include_pad": count_include_pad,
+        "divisor_override": divisor_override,
         "input_shape": input_shape,
         "output_shape": output_shape
     }))
@@ -377,7 +400,13 @@ def _avgpool2d_to_linear_matrix(
     output_shape: Tuple[int, ...],
     kernel_size: int,
     stride: int,
-    padding: int
+    padding: int,
+    *,
+    ceil_mode: bool = False,
+    count_include_pad: bool = True,
+    divisor_override=None,
+    device=None,
+    dtype=None,
 ) -> torch.Tensor:
     """Convert AvgPool2d to equivalent linear transformation matrix."""
     _, channels, in_h, in_w = input_shape
@@ -386,7 +415,12 @@ def _avgpool2d_to_linear_matrix(
     input_flat_size = channels * in_h * in_w
     output_flat_size = channels * out_h * out_w
 
-    W_equiv = torch.zeros(output_flat_size, input_flat_size)
+    W_equiv = torch.zeros(
+        output_flat_size,
+        input_flat_size,
+        device=device,
+        dtype=dtype or torch.get_default_dtype(),
+    )
 
     if isinstance(kernel_size, int):
         kernel_size = (kernel_size, kernel_size)
@@ -398,22 +432,31 @@ def _avgpool2d_to_linear_matrix(
     kernel_h, kernel_w = kernel_size
 
     c, out_y, out_x, k_y, k_x = torch.meshgrid(
-        torch.arange(channels),
-        torch.arange(out_h),
-        torch.arange(out_w),
-        torch.arange(kernel_h),
-        torch.arange(kernel_w),
+        torch.arange(channels, device=W_equiv.device),
+        torch.arange(out_h, device=W_equiv.device),
+        torch.arange(out_w, device=W_equiv.device),
+        torch.arange(kernel_h, device=W_equiv.device),
+        torch.arange(kernel_w, device=W_equiv.device),
         indexing="ij",
     )
 
     in_y = out_y * stride[0] - padding[0] + k_y
     in_x = out_x * stride[1] - padding[1] + k_x
     valid = (in_y >= 0) & (in_y < in_h) & (in_x >= 0) & (in_x < in_w)
 
-    valid_count = valid.sum(dim=(-2, -1), keepdim=True)
-    weight_vals = torch.zeros_like(valid_count, dtype=W_equiv.dtype)
-    nonzero = valid_count > 0
-    weight_vals[nonzero] = valid_count[nonzero].to(W_equiv.dtype).reciprocal()
+    denominators = avgpool2d_denominators(
+        (in_h, in_w),
+        (out_h, out_w),
+        kernel_size,
+        stride,
+        padding,
+        ceil_mode=ceil_mode,
+        count_include_pad=count_include_pad,
+        divisor_override=divisor_override,
+        device=W_equiv.device,
+        dtype=W_equiv.dtype,
+    )
+    weight_vals = denominators.reciprocal().view(1, out_h, out_w, 1, 1)
 
     out_idx = (c * (out_h * out_w) + out_y * out_w + out_x)[valid]
     in_idx = (c * (in_h * in_w) + in_y * in_w + in_x)[valid]

act/back_end/layer_schema.py

@@ -426,6 +426,7 @@ class LayerKind(str, enum.Enum):
             "dilation",
             "ceil_mode",
             "count_include_pad",
+            "divisor_override",
             "output_size",
             "input_shape",
             "output_shape",

act/back_end/net_factory.py

@@ -2174,7 +2174,8 @@ def _lt_spec_cnn_pool() -> Dict[str, Any]:
     # hybridz_tf/tf_cnn.py:44-86 is unreachable via the random CNN generator
     # because it doesn't emit MAXPOOL2D as a direct child of CONV2D inside a
     # hz_cache-bearing context. AvgPool2D additionally exercises the average
-    # pool interval branch in interval_tf/tf_cnn.py.
+    # pool interval branch in interval_tf/tf_cnn.py. The two padded average
+    # pools cover include-pad with a ceil-mode tail and exclude-pad semantics.
     return {"layers": _lt_input([1, 1, 8, 8], -1.0, 1.0) + [
         {"kind": LayerKind.CONV2D.value, "params": {
             "in_channels": 1, "out_channels": 2, "kernel_size": 3,
@@ -2186,8 +2187,14 @@ def _lt_spec_cnn_pool() -> Dict[str, Any]:
             "input_shape": [1, 2, 8, 8], "output_shape": [1, 2, 4, 4],
         }},
         {"kind": LayerKind.AVGPOOL2D.value, "params": {
-            "kernel_size": 2, "stride": 2, "padding": 0,
-            "input_shape": [1, 2, 4, 4], "output_shape": [1, 2, 2, 2],
+            "kernel_size": [3, 2], "stride": [2, 2], "padding": [1, 1],
+            "ceil_mode": True, "count_include_pad": True,
+            "input_shape": [1, 2, 4, 4], "output_shape": [1, 2, 3, 3],
+        }},
+        {"kind": LayerKind.AVGPOOL2D.value, "params": {
+            "kernel_size": 2, "stride": 2, "padding": 1,
+            "ceil_mode": False, "count_include_pad": False,
+            "input_shape": [1, 2, 3, 3], "output_shape": [1, 2, 2, 2],
         }},
         {"kind": LayerKind.FLATTEN.value, "params": {"start_dim": 1}},
         _lt_assert_le([1.0] + [0.0] * 7, 100.0),
@@ -2509,4 +2516,4 @@ def _lt_spec_scale() -> Dict[str, Any]:
     "build_mlp_layers",
     "build_cnn_layers",
     "LAYER_TESTING_SPECS",
-]
+]