tensor-ee 0.1.0__py3-none-any.whl

tensor_ee/blocks/raster/conv.py

@@ -0,0 +1,183 @@
+"""R convolution: PyTorch ``nn.Conv2d`` as ordered neighbourhood bands and
+``matrixMultiply``.
+
+Moved whole from ``tensor_ee.compiler`` (``conv2d``, ``_pair``, ``_names``) and
+``tensor_ee.semantics`` (``conv_options`` with the coefficient checks it uses:
+``_coefficient_field``, ``finite_coefficients``). The old names are re-exports
+of these objects. The errors they raise, ``SemanticUnknown`` and ``reject``
+(here ``_reject``), are in ``tensor_ee.core.errors``.
+
+Supported settings (contract section 8): odd kernels with same padding, stride
+1 or 2, dilation 1, groups 1, zero padding, finite float32 coefficients. Other
+strides stay rejected.
+
+Precision policy ``float32_operands`` (contract section 7): the weight and bias
+arrays are cast ``.float()`` before ``matrixMultiply``. The node output is
+materialized with ``.toFloat()`` by the caller (``compile_model`` or the
+``tensor_ee.primitives`` facade), not here.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+import ee
+import torch
+from torch import nn
+
+from ...core.errors import SemanticUnknown, UnsupportedOperation, reject as _reject
+from ...core.metadata import ANY_SETTINGS, BlockMetadata, Locality, validate_metadata
+from ...core.precision import FLOAT32_OPERANDS, validate_policy
+from ...core.tensor import Raster
+# Library converters and coefficient transport. These module-level bindings are
+# the names conv2d looks up (tests patch _inline_float32_array here).
+from ...layouts.raster import project as _project, stride_sample as _stride_sample
+from ...layouts.raster import zero_extended as _zero_extended
+from ...weights.constants import inline_float32_array as _inline_float32_array
+
+
+def _pair(value: Any) -> tuple[Any, Any]:
+    if isinstance(value, (tuple, list)):
+        if len(value) != 2:
+            raise UnsupportedOperation(f"Expected a scalar or pair, received {value!r}")
+        return tuple(value)
+    return value, value
+
+
+def _names(count: int) -> tuple[str, ...]:
+    return tuple(f"c{i}" for i in range(count))
+
+
+def finite_coefficients(module: nn.Module) -> None:
+    """Inspect local coefficients only, never an input-shaped activation."""
+    for name, value in (*module.named_parameters(recurse=False), *module.named_buffers(recurse=False)):
+        if not value.is_floating_point():
+            continue
+        if value.dtype != torch.float32:
+            _reject("Only float32 model parameters and buffers are supported; convert with model.float()", "model_dtype")
+        try:
+            finite = bool(torch.isfinite(value.detach()).all().item())
+        except Exception as error:
+            raise SemanticUnknown(f"Cannot inspect local coefficient {name}: {type(error).__name__}: {error}") from error
+        if not finite:
+            _reject(f"Local coefficient {name} must contain only finite tensors", "nonfinite_coefficients")
+
+
+def _coefficient_field(value: Any, name: str, shape: tuple[int, ...]) -> None:
+    # Computational tensors require float32 even though the report's global
+    # floating-dtype Boolean deliberately ignores integer bookkeeping.
+    if not isinstance(value, torch.Tensor) or value.dtype != torch.float32:
+        _reject(f"Only float32 computational coefficient {name} is supported", "model_dtype")
+    try:
+        actual = tuple(value.shape)
+    except Exception as error:
+        raise SemanticUnknown(f"Cannot inspect local coefficient {name} shape: {error}") from error
+    if actual != shape:
+        _reject(f"Local coefficient {name} shape must be {shape}, received {actual}", "coefficient_shape")
+
+
+def conv_options(module: nn.Conv2d) -> tuple[tuple[int, int], int]:
+    kh, kw = _pair(module.kernel_size)
+    if kh % 2 != 1 or kw % 2 != 1:
+        _reject("Conv2d kernels must have odd height and width")
+    padding = (kh // 2, kw // 2) if module.padding == "same" else module.padding
+    if padding != (kh // 2, kw // 2):
+        _reject("Conv2d must use same padding (kernel_size // 2)")
+    stride = _pair(module.stride)
+    if (stride not in ((1, 1), (2, 2)) or _pair(module.dilation) != (1, 1)
+            or module.groups != 1 or module.padding_mode != "zeros"):
+        _reject("Conv2d requires stride=1 or 2, dilation=1, groups=1, padding_mode='zeros'")
+    _coefficient_field(module.weight, "Conv2d.weight", (module.out_channels, module.in_channels, kh, kw))
+    if module.bias is not None:
+        _coefficient_field(module.bias, "Conv2d.bias", (module.out_channels,))
+    finite_coefficients(module)
+    return (kh, kw), stride[0]
+
+
+def conv2d(raster: Raster, module: nn.Conv2d,
+           parameters: tuple[ee.Array, ee.Array | None] | None = None,
+           *, execution_dtype: str = "float32", clip_footprints: bool = False,
+           policy: str = FLOAT32_OPERANDS) -> Raster:
+    """PyTorch cross-correlation via ordered neighborhood bands and matrixMultiply."""
+    validate_policy(policy, allowed=FLOAT32_OPERANDS)
+    if execution_dtype != "float32":
+        raise ValueError("execution_dtype must be float32 (maximum supported precision)")
+    (kh, kw), stride_value = conv_options(module)
+    stride = (stride_value, stride_value)
+    if module.in_channels != len(raster.bands):
+        raise UnsupportedOperation("Conv2d input channel count does not match raster")
+    kernel = ee.Kernel.fixed(kw, kh, [[1] * kw for _ in range(kh)], kw // 2, kh // 2, False)
+    # Never rely on the EE-generated default band ordering. Torch's flattened
+    # OIHW weights require channels, then rows, then columns. EE suffixes use
+    # offsets relative to the kernel origin, with x before y.
+    patch_bands = [f"{band}_{dx}_{dy}" for band in raster.bands
+                   for dy in range(-(kh // 2), kh // 2 + 1)
+                   for dx in range(-(kw // 2), kw // 2 + 1)]
+    patch = (_zero_extended(raster).neighborhoodToBands(kernel)
+             .select(patch_bands).toArray().toArray(1))
+    if parameters is None:
+        weights = _inline_float32_array(module.weight.detach().cpu().float()
+                           .reshape(module.out_channels, -1).tolist())
+        bias = (None if module.bias is None else
+                _inline_float32_array([[value] for value in module.bias.detach().cpu().float().tolist()]))
+    else:
+        weights, bias = parameters
+    if execution_dtype == "float32":
+        # Keep matrix operands in float32 rather than promoting the entire
+        # activation neighborhood to double through the frozen weights.
+        weights = weights.float()
+        if bias is not None:
+            bias = bias.float()
+    result = ee.Image(weights).matrixMultiply(patch)
+    if bias is not None:
+        result = result.add(ee.Image(bias))
+    bands = _names(module.out_channels)
+    result = result.arrayProject([0]).arrayFlatten([list(bands)])
+    # Make the dense convolution's sampling grid explicit before decimation.
+    # Stride one already receives that projection in _finish; avoid a second
+    # redundant reproject in the long residual/decoder graph.
+    if stride[0] > 1:
+        result = _project(result, raster.grid, clip_footprints)
+    return _stride_sample(result, bands, raster.grid, stride[0], clip_footprints)
+
+
+# ---------------------------------------------------------------------------
+# Locality and cost metadata (contract section 9)
+# ---------------------------------------------------------------------------
+def _conv_coefficients(module: nn.Conv2d) -> int:
+    """Values embedded in the request when the coefficients are inline:
+    ``out * in * kh * kw`` weights plus ``out`` biases."""
+    kh, kw = _pair(module.kernel_size)
+    count = module.out_channels * module.in_channels * kh * kw
+    return count + (module.out_channels if module.bias is not None else 0)
+
+
+def _conv_asset_arrays(module: nn.Conv2d) -> int:
+    """Server arrays read when the coefficients are supplied as EE arrays:
+    the weight matrix, plus the bias column."""
+    return 1 + (module.bias is not None)
+
+
+def _no_coefficients(*args, **kwargs) -> int:
+    """Cost of an R block without weights: it embeds no coefficients and reads
+    no asset arrays. Scalar options (a slope, a threshold, the float32 maximum)
+    are graph literals, not coefficients, and are not counted."""
+    return 0
+
+
+METADATA = validate_metadata({
+    # Radius k // 2 for a k x k kernel (receptive_field._window_axis with
+    # padding k // 2). The static field records the largest validated kernel,
+    # 7 (radius 3); kernels 1 and 3 have radius 0 and 1.
+    # Validated: any accepted call. Kernels 1, 3 and 7 at strides 1 and 2, and
+    # clip_footprints both ways (True in resnet18_full64_float32_policy.json), were run;
+    # the unqualified "Ready" covers them, and other forms are refused.
+    "conv2d": BlockMetadata(
+        name="conv2d", layout_in="R", layout_out="R", locality=Locality.KERNEL,
+        policy=FLOAT32_OPERANDS, embedded_constants=_conv_coefficients,
+        asset_arrays=_conv_asset_arrays, radius=3, status="Ready",
+        evidence=("docs/history/cnn/results/parity.json", "docs/history/cnn/results/primitives_float32.json",
+                  "docs/history/cnn/results/resnet_primitives.json", "docs/history/cnn/results/resnet18_full64_float32_policy.json",
+                  "docs/history/cnn/results/segwater_js_parity_1x1.json", "docs/history/cnn/results/tiled_resnet18_2x2_parity.json"),
+        validated=ANY_SETTINGS),
+})

tensor_ee/blocks/raster/merge.py

@@ -0,0 +1,66 @@
+"""R merges: channel concatenation (``torch.cat`` on dim 1) and residual addition.
+
+Moved whole from ``tensor_ee.compiler`` (``concatenate``, ``residual_add``).
+The old names are re-exports of these objects. Both require identical grids;
+there is no broadcasting.
+"""
+
+from __future__ import annotations
+
+from typing import Sequence
+
+from ...core.errors import UnsupportedOperation
+from ...core.metadata import ANY_SETTINGS, BlockMetadata, Locality, Validated, validate_metadata
+from ...core.precision import FLOAT32_OPERANDS, validate_policy
+from ...core.tensor import Raster
+from ...layouts.raster import finish as _finish
+from .conv import _names, _no_coefficients
+
+
+def concatenate(rasters: Sequence[Raster], dim: int = 1,
+                 *, clip_footprints: bool = False, policy: str = FLOAT32_OPERANDS) -> Raster:
+    validate_policy(policy, allowed=FLOAT32_OPERANDS)
+    if dim not in (1, -3) or not rasters:
+        raise UnsupportedOperation("Concatenation requires a nonempty sequence and NCHW channel dim=1")
+    grid = rasters[0].grid
+    if any(raster.grid != grid for raster in rasters):
+        raise UnsupportedOperation("Skip connection tensors must have identical spatial grids")
+    parts = [raster.image.rename([f"part{i}_{band}" for band in raster.bands])
+             for i, raster in enumerate(rasters)]
+    result = parts[0]
+    for part in parts[1:]:
+        result = result.addBands(part)
+    return _finish(result, _names(sum(len(raster.bands) for raster in rasters)), grid, clip_footprints)
+
+
+def residual_add(left: Raster, right: Raster, *, clip_footprints: bool = False,
+                 policy: str = FLOAT32_OPERANDS) -> Raster:
+    """Elementwise addition of equally shaped tensors; broadcasting is rejected."""
+    validate_policy(policy, allowed=FLOAT32_OPERANDS)
+    if not isinstance(left, Raster) or not isinstance(right, Raster):
+        raise UnsupportedOperation("Residual addition requires two raster tensors")
+    if left.grid != right.grid or len(left.bands) != len(right.bands):
+        raise UnsupportedOperation("Residual addition requires identical spatial grids and channel counts")
+    return _finish(left.image.add(right.image.rename(list(left.bands))), left.bands, left.grid, clip_footprints)
+
+
+# ---------------------------------------------------------------------------
+# Locality and cost metadata (contract section 9)
+# ---------------------------------------------------------------------------
+METADATA = validate_metadata({
+    # The U-Net's skip concatenations are torch.cat(dim=1); dim -3, the same axis, never ran.
+    "concatenate": BlockMetadata(
+        name="concatenate", layout_in="R", layout_out="R", locality=Locality.POINTWISE,
+        policy=FLOAT32_OPERANDS, embedded_constants=_no_coefficients,
+        asset_arrays=_no_coefficients, status="Ready",
+        evidence=("docs/history/cnn/results/parity.json", "docs/history/cnn/results/resnet18_full64_float32_policy.json",
+                  "docs/history/cnn/results/segwater_js_parity_1x1.json"),
+        validated=Validated(({"dim": (1,)},))),
+    "residual_add": BlockMetadata(
+        name="residual_add", layout_in="R", layout_out="R", locality=Locality.POINTWISE,
+        policy=FLOAT32_OPERANDS, embedded_constants=_no_coefficients,
+        asset_arrays=_no_coefficients, status="Ready",
+        evidence=("docs/history/cnn/results/resnet_primitives.json", "docs/history/cnn/results/resnet18_full64_float32_policy.json",
+                  "docs/history/cnn/results/segwater_js_parity_1x1.json"),
+        validated=ANY_SETTINGS),
+})

tensor_ee/blocks/raster/outputs.py

@@ -0,0 +1,211 @@
+"""Recipe-free segmentation outputs on explicit finite raster grids.
+
+Scores/probabilities are float32. Classes are the ordered raw-logit argmax;
+thresholds compare widened probabilities against a retained binary64 scalar.
+These functions build deferred graphs without cloud inspection or initialization.
+
+Moved whole from ``tensor_ee.outputs`` (``segmentation_outputs``,
+``binary_threshold``, ``SegmentationOutputs`` and their helpers). The old names
+are re-exports of these objects. The probabilities are this module's softmax
+block materialized with ``.toFloat()``, exactly as the public
+``tensor_ee.softmax`` primitive the old module called.
+
+Precision policy ``float32_operands`` (contract section 7): the threshold
+widens the float32 probabilities ``toDouble()`` and compares them with an exact
+binary64 scalar (``_ee_scalar``). The generic defaults (positive class 1,
+threshold 0.5) are API defaults, not the SegWater recipe (contract section 11).
+"""
+from __future__ import annotations
+
+from dataclasses import dataclass
+from numbers import Integral, Real
+import struct
+from typing import Sequence
+
+import ee
+
+from ...core.errors import UnsupportedOperation
+from ...core.metadata import ANY_SETTINGS, BlockMetadata, Locality, validate_metadata
+from ...core.precision import FLOAT32_OPERANDS, validate_policy
+from ...core.precision import exact_scalar as _ee_scalar
+from ...core.tensor import Raster
+from ...layouts.raster import _float32, finish as _finish
+from .activations import _real_option, common_score_mask as _common_score_mask, softmax as _softmax
+from .conv import _no_coefficients
+
+
+@dataclass(frozen=True)
+class SegmentationOutputs:
+    logits: Raster
+    probabilities: Raster
+    classes: Raster
+
+
+def _band_names(names: Sequence[str], count: int, option: str) -> tuple[str, ...]:
+    if isinstance(names, str):
+        raise UnsupportedOperation(f"{option} must be a sequence of {count} names")
+    try:
+        result = tuple(names)
+    except TypeError as error:
+        raise UnsupportedOperation(f"{option} must be a sequence of {count} names") from error
+    if (len(result) != count or any(not isinstance(name, str) or not name for name in result)
+            or len(set(result)) != count):
+        raise UnsupportedOperation(f"{option} must contain exactly {count} nonempty unique names")
+    return result
+
+
+def _scores(raster: Raster, *, binary: bool = False) -> ee.Image:
+    if not isinstance(raster, Raster):
+        raise TypeError("scores must be a Raster")
+    count = len(raster.bands)
+    if count < 2 or (binary and count != 2):
+        requirement = "exactly two" if binary else "at least two"
+        raise UnsupportedOperation(f"scores require {requirement} channels")
+    _band_names(raster.bands, count, "score bands")
+    return raster.image.select(list(raster.bands))
+
+
+def _check_literal_validity(image: ee.Image) -> None:
+    """Inspect only constant literals and rename/cast/reproject wrappers.
+
+    No generic graph evaluator is attempted. Computed literals, arbitrary graph
+    operations and remote assets remain the caller's binary/finiteness assertion.
+    A float cast is applied to the literal values before checking the effective
+    binary payload; masks/other operations stop this narrow inspection.
+    """
+    wrappers = {"Image.rename": "input", "Image.toFloat": "value",
+                "Image.toDouble": "value", "Image.reproject": "image"}
+    float_cast = False
+    while isinstance(image, ee.Image) and image.func is not None:
+        signature = image.func.getSignature()
+        name = signature.get("name")
+        if name not in wrappers:
+            break
+        float_cast = float_cast or name == "Image.toFloat"
+        image = image.args[wrappers[name]]
+    if not isinstance(image, ee.Image) or image.func is None:
+        return
+    if image.func.getSignature().get("name") != "Image.constant":
+        return
+    value = image.args.get("value")
+    if isinstance(value, Real):
+        values = (value,)
+    elif isinstance(value, (list, tuple)) and all(isinstance(v, Real) for v in value):
+        values = tuple(value)
+    else:
+        return
+    if len(values) != 1:
+        raise UnsupportedOperation("local validity must have one band matching its Raster metadata")
+    if float_cast:
+        try:
+            values = tuple(struct.unpack("f", struct.pack("f", float(v)))[0] for v in values)
+        except (OverflowError, ValueError) as error:
+            raise UnsupportedOperation("local validity must contain finite binary values") from error
+    if any(v not in (0, 1) for v in values):
+        raise UnsupportedOperation("local validity must contain finite binary values")
+
+
+def _validity_mask(validity: Raster | None, raster: Raster) -> ee.Image | None:
+    if validity is None:
+        return None
+    if not isinstance(validity, Raster):
+        raise TypeError("validity must be a Raster")
+    if validity.grid != raster.grid:
+        raise UnsupportedOperation("validity must use the identical output grid")
+    _band_names(validity.bands, 1, "validity bands")
+    _check_literal_validity(validity.image)
+    image = validity.image.select(list(validity.bands))
+    return image.eq(1).And(image.mask().gt(0)).unmask(0, sameFootprint=False)
+
+
+def _masked_scores(raster: Raster, image: ee.Image, validity: Raster | None) -> Raster:
+    mask = _common_score_mask(image, raster.grid)
+    explicit = _validity_mask(validity, raster)
+    if explicit is not None:
+        mask = mask.And(explicit)
+    return Raster(image.updateMask(mask), raster.bands, raster.grid)
+
+
+def segmentation_outputs(logits: Raster, *, logit_bands: Sequence[str] | None = None,
+                         probability_bands: Sequence[str] | None = None,
+                         class_band: str = "class_index",
+                         validity: Raster | None = None,
+                         policy: str = FLOAT32_OPERANDS) -> SegmentationOutputs:
+    """Retain raw logits and return float32 probabilities/int32 logit argmax.
+
+    Ordered channel indices define classes; exact logit ties choose the first.
+    Processed outputs share all score masks/finiteness, finite grid and optional
+    explicit validity. Raw logits retain their original values and masks.
+    """
+    validate_policy(policy, allowed=FLOAT32_OPERANDS)
+    image = _scores(logits)
+    count = len(logits.bands)
+    logit_names = logits.bands if logit_bands is None else _band_names(logit_bands, count, "logit_bands")
+    probability_names = (_band_names(probability_bands, count, "probability_bands") if probability_bands is not None
+                         else tuple(f"probability_{i}" for i in range(count)))
+    class_name = _band_names((class_band,), 1, "class_band")
+    raw = logits if logit_bands is None else Raster(image.rename(list(logit_names)), logit_names, logits.grid)
+    masked = _masked_scores(logits, image, validity)
+    # The softmax block materialized as float32, as tensor_ee.softmax does.
+    probabilities = _float32(_softmax(masked))
+    probabilities = Raster(probabilities.image.rename(list(probability_names)), probability_names, logits.grid)
+    # Do not derive the winner from rounded probabilities: [0,1e-8] can round
+    # to [.5,.5]. arrayArgmax returns the first maximum of the ordered logits.
+    winner = masked.image.toArray().arrayArgmax().arrayGet([0]).toInt32()
+    classes = _finish(winner, class_name, logits.grid)
+    return SegmentationOutputs(raw, probabilities, classes)
+
+
+def _threshold_option(threshold: float) -> float:
+    if isinstance(threshold, bool) or not isinstance(threshold, Real):
+        raise UnsupportedOperation("threshold must be a static finite real scalar")
+    if not 0 <= threshold <= 1:
+        raise UnsupportedOperation("threshold must satisfy 0 <= threshold <= 1")
+    number = _real_option(threshold, "threshold")
+    if threshold != 0 and number == 0:
+        raise UnsupportedOperation("threshold must be representable in binary64 without underflow to zero")
+    return number
+
+
+def binary_threshold(probabilities: Raster, *, positive_class: int = 1,
+                     threshold: float = 0.5, output_band: str = "class_mask",
+                     validity: Raster | None = None,
+                     policy: str = FLOAT32_OPERANDS) -> Raster:
+    """Inclusive uint8 threshold on one of exactly two normalized channels.
+
+    Widen float32 probabilities losslessly and retain the finite binary64
+    threshold, so probability .5 is below threshold .50000001. Invalid pixels
+    are excluded rather than filled with class zero. Caller supplies normalized
+    probabilities and asserts binary/finiteness for computed/remote validity.
+    """
+    validate_policy(policy, allowed=FLOAT32_OPERANDS)
+    image = _scores(probabilities, binary=True)
+    if isinstance(positive_class, bool) or not isinstance(positive_class, Integral) or positive_class not in (0, 1):
+        raise UnsupportedOperation("positive_class must be the static index 0 or 1")
+    value = _threshold_option(threshold)
+    name = _band_names((output_band,), 1, "output_band")
+    masked = _masked_scores(probabilities, image, validity)
+    decision = masked.image.select(probabilities.bands[int(positive_class)]).toDouble().gte(_ee_scalar(value)).toUint8()
+    return _finish(decision, name, probabilities.grid)
+
+
+# ---------------------------------------------------------------------------
+# Locality and cost metadata (contract section 9)
+# ---------------------------------------------------------------------------
+METADATA = validate_metadata({
+    # Logits, float32 probabilities and an int32 class band, all on the input grid.
+    # Validated: any accepted call (segmentation output validation varied the validity input; band names are labels).
+    "segmentation_outputs": BlockMetadata(
+        name="segmentation_outputs", layout_in="R", layout_out="R", locality=Locality.POINTWISE,
+        policy=FLOAT32_OPERANDS, embedded_constants=_no_coefficients,
+        asset_arrays=_no_coefficients, status="Ready",
+        evidence=("docs/evidence/t03-outputs.json", "docs/evidence/t05-primary.json"),
+        validated=ANY_SETTINGS),
+    # Validated: any accepted call (segmentation output validation: thresholds 0, 0.5, 1 and 0.50000001, both classes).
+    "binary_threshold": BlockMetadata(
+        name="binary_threshold", layout_in="R", layout_out="R", locality=Locality.POINTWISE,
+        policy=FLOAT32_OPERANDS, embedded_constants=_no_coefficients,
+        asset_arrays=_no_coefficients, status="Ready",
+        evidence=("docs/evidence/t03-outputs.json",),
+        validated=ANY_SETTINGS),
+})

tensor_ee/blocks/raster/pooling.py

@@ -0,0 +1,182 @@
+"""R pooling: PyTorch ``nn.AvgPool2d`` and ``nn.MaxPool2d`` at a small set of
+common CNN configurations.
+
+Moved whole from ``tensor_ee.compiler`` (``pool2d``, both branches) and
+``tensor_ee.semantics`` (``pool_options``). The old names are re-exports of
+these objects. The raster fixtures added the configurations beyond the two the reference
+models use (``pool_config``).
+
+Supported settings (contract section 8), all with ``ceil_mode=False``,
+max-pool ``dilation=1`` and ``return_indices=False``, and average-pool
+``divisor_override=None``:
+
+- **aligned** average or max, kernel = stride = k, padding 0, for k in
+  (2, 3, 4), on input sizes divisible by k: ``reduceResolution`` onto the
+  k-times coarser grid;
+- **max 3 x 3, padding 1, stride 1 or 2**: a 3 x 3 neighbourhood maximum on
+  the input grid (then stride sampling for stride 2). The input is masked
+  outside its grid and Earth Engine reducers skip masked pixels, so the padding
+  never wins: PyTorch's negative-infinity padding;
+- **average 3 x 3, stride 1, padding 1**, under both PyTorch divisor rules:
+  ``count_include_pad=True`` divides the in-bounds sum by 9 (zero padding);
+  ``count_include_pad=False`` takes the mean of the in-bounds pixels only.
+
+Every other configuration is rejected with the list above. Each supported
+configuration has a raster-setting validation fixture (``raster/pool2d/*`` in
+``tools/fixtures/raster.py``) that compares it with PyTorch: bit for bit for
+max pooling, within 1e-4 for averages.
+
+Precision policy ``float32_operands`` (contract section 7). There are no
+coefficients. The node output is materialized with ``.toFloat()`` by the caller
+(``compile_model`` or the ``tensor_ee.primitives`` facade).
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import ee
+from torch import nn
+
+from ...core.errors import UnsupportedOperation
+from ...core.metadata import BlockMetadata, Locality, Validated, validate_metadata
+from ...core.precision import FLOAT32_OPERANDS, validate_policy
+from ...core.tensor import Raster
+from ...layouts.raster import finish as _finish, project as _project
+from ...layouts.raster import stride_sample as _stride_sample
+from ...core.errors import reject as _reject
+from .conv import _no_coefficients, _pair
+
+ALIGNED_KERNELS = (2, 3, 4)
+"""Kernel sizes k of the aligned k x k pools (stride k, padding 0)."""
+
+SUPPORTED = ("aligned avg/max k x k with stride k and padding 0 for k in (2, 3, 4) on sizes divisible "
+             "by k; max 3x3 padding 1 with stride 1 or 2; avg 3x3 stride 1 padding 1")
+"""The supported configurations, as the rejection message names them."""
+
+
+@dataclass(frozen=True)
+class PoolConfig:
+    """One supported pooling configuration, normalized from the PyTorch module.
+
+    ``reducer`` is ``"max"`` or ``"mean"``; ``kernel``, ``stride`` and
+    ``padding`` are the equal per-axis values; ``count_include_pad`` is the
+    average pool's divisor rule (``None`` for max pooling).
+    """
+
+    reducer: str
+    kernel: int
+    stride: int
+    padding: int
+    count_include_pad: bool | None = None
+
+    @property
+    def aligned(self) -> bool:
+        """An aligned k x k pool: kernel = stride, no padding."""
+        return self.padding == 0
+
+    def output_size(self, size: int) -> int:
+        """PyTorch's output size on one axis (``ceil_mode=False``)."""
+        return (size + 2 * self.padding - self.kernel) // self.stride + 1
+
+
+def aligned_size_message(kernel: int) -> str:
+    """The rejection of an input size that an aligned k x k pool does not divide."""
+    if kernel == 2:
+        return "This experiment requires even spatial dimensions before pooling"
+    return f"Aligned {kernel}x{kernel} pooling requires spatial dimensions divisible by {kernel}"
+
+
+def pool_config(module: nn.AvgPool2d | nn.MaxPool2d) -> PoolConfig:
+    """Validate the module's options and return its supported configuration."""
+    stride = module.stride if module.stride is not None else module.kernel_size
+    if module.ceil_mode is not False:
+        _reject("Pooling requires ceil_mode=False")
+    if isinstance(module, nn.MaxPool2d):
+        if _pair(module.dilation) != (1, 1) or module.return_indices is not False:
+            _reject("MaxPool2d requires dilation=1 and return_indices=False")
+        reducer, include = "max", None
+    elif module.divisor_override is not None or not isinstance(module.count_include_pad, bool):
+        _reject("AvgPool2d requires divisor_override=None and static count_include_pad")
+    else:
+        reducer, include = "mean", module.count_include_pad
+    kernel, strides, padding = _pair(module.kernel_size), _pair(stride), _pair(module.padding)
+    for k in ALIGNED_KERNELS:
+        if kernel == (k, k) and strides == (k, k) and padding == (0, 0):
+            return PoolConfig(reducer, k, k, 0, include)
+    if kernel == (3, 3) and padding == (1, 1):
+        if strides == (1, 1) or (strides == (2, 2) and reducer == "max"):
+            return PoolConfig(reducer, 3, int(strides[0]), 1, include)
+    _reject(f"Pooling supports {SUPPORTED}, without ceil_mode; received {type(module).__name__}"
+            f"(kernel_size={module.kernel_size!r}, stride={stride!r}, padding={module.padding!r})")
+
+
+def pool_options(module: nn.AvgPool2d | nn.MaxPool2d) -> bool:
+    """Validate the module and return whether it pools a padded 3x3 neighbourhood
+    (``pool_config(module).padding``), not an aligned k x k block."""
+    return not pool_config(module).aligned
+
+
+def pool2d(raster: Raster, module: nn.AvgPool2d | nn.MaxPool2d,
+           *, clip_footprints: bool = False, policy: str = FLOAT32_OPERANDS) -> Raster:
+    """PyTorch average or max pooling at a supported configuration (``pool_config``)."""
+    validate_policy(policy, allowed=FLOAT32_OPERANDS)
+    config = pool_config(module)
+    if not config.aligned:
+        kernel = ee.Kernel.square(1, "pixels", False)
+        if config.reducer == "max":
+            # Masked padding contributes no value to max, just as -infinity.
+            result = raster.image.reduceNeighborhood(ee.Reducer.max(), kernel, skipMasked=False)
+        elif config.count_include_pad:
+            # count_include_pad=True: the in-bounds sum over the full window of 9
+            # (PyTorch's zero padding counts in the divisor).
+            result = raster.image.reduceNeighborhood(ee.Reducer.sum(), kernel, skipMasked=False).divide(9)
+        else:
+            # count_include_pad=False: the mean of the in-bounds pixels only.
+            result = raster.image.reduceNeighborhood(ee.Reducer.mean(), kernel, skipMasked=False)
+        if config.stride == 1:
+            # The output is on the input grid; _stride_sample(..., 1) is _finish,
+            # which pins it there (as conv2d does at stride 1).
+            return _stride_sample(result, raster.bands, raster.grid, 1, clip_footprints)
+        result = _project(result, raster.grid, clip_footprints)
+        return _stride_sample(result, raster.bands, raster.grid, config.stride, clip_footprints)
+    k = config.kernel
+    reducer = ee.Reducer.max() if config.reducer == "max" else ee.Reducer.mean()
+    if raster.grid.width % k or raster.grid.height % k:
+        raise UnsupportedOperation(aligned_size_message(k))
+    grid = raster.grid.resized(k, raster.grid.width // k, raster.grid.height // k)
+    # maxPixels (k + 2)^2 keeps today's 16 at k = 2: the k x k block plus a
+    # one-pixel ring. It bounds the input pixels per output pixel; with
+    # bestEffort=False a larger need fails the request rather than changing values.
+    result = raster.image.reduceResolution(reducer, bestEffort=False, maxPixels=(k + 2) ** 2)
+    result = _project(result, grid, clip_footprints)
+    return _finish(result, raster.bands, grid, clip_footprints)
+
+
+# ---------------------------------------------------------------------------
+# Locality and cost metadata (contract section 9)
+# ---------------------------------------------------------------------------
+METADATA = validate_metadata({
+    # receptive_field._window_axis: output pixel i of an aligned k x k pool reads
+    # input pixels [k i, k i + k - 1], and a padded 3x3 pool with stride s reads
+    # [s i - 1, s i + 1]. The radius bounds the offset from s i: k - 1 = 3 at the
+    # largest aligned kernel, 4; 1 for the padded pools.
+    # Validated, in PoolConfig's terms (``pool_config(module)``; "mean" is AvgPool2d): the
+    # status's three families. count_include_pad does not change an unpadded pool, so the
+    # aligned pools do not name it.
+    "pool2d": BlockMetadata(
+        name="pool2d", layout_in="R", layout_out="R", locality=Locality.KERNEL,
+        policy=FLOAT32_OPERANDS, embedded_constants=_no_coefficients,
+        asset_arrays=_no_coefficients, radius=3,
+        status=("Ready — aligned avg/max k = 2, 3, 4; max 3×3 stride 1 or 2, padding 1; "
+                "avg 3×3 stride 1 padding 1, both divisor rules"),
+        evidence=("docs/evidence/live-fixtures/20261013T065257Z-p02.json", "docs/evidence/p02-raster-offline.json",
+                  "docs/history/cnn/results/parity.json", "docs/evidence/t05-primary.json",
+                  "docs/history/cnn/results/resnet_primitives.json", "docs/history/cnn/results/resnet18_full64_float32_policy.json",
+                  "docs/history/cnn/results/segwater_js_parity_1x1.json"),
+        validated=Validated((
+            *({"reducer": ("mean", "max"), "kernel": (k,), "stride": (k,), "padding": (0,)} for k in ALIGNED_KERNELS),
+            {"reducer": ("max",), "kernel": (3,), "stride": (1, 2), "padding": (1,)},
+            {"reducer": ("mean",), "kernel": (3,), "stride": (1,), "padding": (1,), "count_include_pad": (True, False)},
+        ), sources=(PoolConfig,))),
+})