stackformers 3.8.2__tar.gz → 3.9.0__tar.gz

{stackformers-3.8.2 → stackformers-3.9.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: stackformers
-Version: 3.8.2
+Version: 3.9.0
 Summary: Typed, composable, SOLID transformer library for PyTorch
 Project-URL: Homepage, https://github.com/Red-Eyed/stackformers
 Project-URL: Repository, https://github.com/Red-Eyed/stackformers

{stackformers-3.8.2 → stackformers-3.9.0}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "stackformers"
-version = "3.8.2"
+version = "3.9.0"
 description = "Typed, composable, SOLID transformer library for PyTorch"
 readme = "README.md"
 license = { file = "LICENSE" }

{stackformers-3.8.2 → stackformers-3.9.0}/stackformers/presets/README.md

@@ -19,7 +19,9 @@ Each preset is split into two levels:
 
 **`CrossAttender`** — cross-attn → ff per layer. `x` is a fixed set of queries with no self-attention. Useful for Perceiver-style or slot-attention architectures.
 
-All presets accept `PaddedInput` or `PackedInput` interchangeably.
+**`PyramidVisionBackbone`** — multi-scale vision backbone (PVTv2-style hybrid). Stages of `spatial/` attention joined by `PatchMerging` downsamples; consumes a `SpatialInput` from your own patch-embed stem and returns one `(b, d, H, W)` feature map per stage. `pyramid_vision_config` defaults to 2-D window attention in the high-resolution early stages and spatially-reduced global attention later, tuned for 1024² images with an 8×8 patch stem. Unlike the other presets it is grid-based, so it takes `SpatialInput` rather than `PaddedInput`/`PackedInput`.
+
+The other presets accept `PaddedInput` or `PackedInput` interchangeably.
 
 ## Customising a collaborator (subclass the concrete preset)
 

stackformers-3.9.0/stackformers/presets/pyramid_vision.py

@@ -0,0 +1,173 @@
+from __future__ import annotations
+
+import torch.nn as nn
+from einops import rearrange
+from jaxtyping import Float
+from pydantic import BaseModel, Field, model_validator
+from torch import Tensor
+
+from stackformers.feedforward.config import FeedForwardConfig, SwiGLUConfig
+from stackformers.feedforward.factory import build_ff
+from stackformers.norm.config import NormConfig, RMSNormConfig
+from stackformers.norm.factory import build_norm
+from stackformers.positional.config import PosEncodingConfig, RoPE2DConfig
+from stackformers.positional.factory import build_pos_encoding
+from stackformers.spatial.config import (
+    PatchMergingConfig,
+    SpatialAttnConfig,
+    SpatialReductionAttentionConfig,
+    WindowAttention2DConfig,
+)
+from stackformers.spatial.factory import build_patch_merging, build_spatial_attn
+from stackformers.spatial.input import SpatialInput
+from stackformers.spatial.layer import SpatialTransformerLayer
+
+
+class PyramidStageConfig(BaseModel):
+    """One pyramid stage: `depth` spatial-attention layers at a fixed grid and width."""
+
+    attn: SpatialAttnConfig
+    ff: FeedForwardConfig
+    norm: NormConfig
+    pos_encoding: PosEncodingConfig
+    depth: int = Field(gt=0)
+
+    @model_validator(mode="after")
+    def _dims_agree(self) -> PyramidStageConfig:
+        if not self.attn.dim == self.ff.dim == self.norm.dim:
+            raise ValueError(
+                f"stage dims disagree: attn={self.attn.dim}, ff={self.ff.dim}, norm={self.norm.dim}"
+            )
+        return self
+
+
+class PyramidVisionConfig(BaseModel):
+    """A multi-scale vision backbone: stages joined by patch-merging downsamples.
+
+    inter_stage[i] merges the output of stage i into the input of stage i+1, so there is
+    exactly one fewer merge than there are stages.
+    """
+
+    stages: list[PyramidStageConfig] = Field(min_length=1)
+    inter_stage: list[PatchMergingConfig]
+
+    @model_validator(mode="after")
+    def _stages_and_merges_chain(self) -> PyramidVisionConfig:
+        if len(self.inter_stage) != len(self.stages) - 1:
+            raise ValueError(
+                f"need {len(self.stages) - 1} inter-stage merges, got {len(self.inter_stage)}"
+            )
+        for i, merge in enumerate(self.inter_stage):
+            if merge.in_dim != self.stages[i].attn.dim:
+                raise ValueError(f"merge {i} in_dim {merge.in_dim} != stage {i} dim")
+            if merge.out_dim != self.stages[i + 1].attn.dim:
+                raise ValueError(f"merge {i} out_dim {merge.out_dim} != stage {i + 1} dim")
+        return self
+
+
+def pyramid_vision_config(
+    *,
+    dims: tuple[int, ...] = (64, 128, 320, 512),
+    depths: tuple[int, ...] = (3, 4, 6, 3),
+    heads: tuple[int, ...] = (1, 2, 5, 8),
+    sr_ratios: tuple[int, ...] = (8, 4, 2, 1),
+    window: int = 8,
+    window_stages: int = 2,
+    ff_mult: float = 4.0,
+    dropout: float = 0.0,
+) -> PyramidVisionConfig:
+    """PVTv2-style hybrid backbone: 2-D window attention in the first `window_stages`
+    high-resolution stages, spatially-reduced global attention in the rest.
+
+    Defaults are tuned for 1024² images patch-embedded with an 8×8 stem (stage-1 grid
+    128×128); `window=8` and the sr_ratios divide that grid cleanly. RoPE 2-D positions
+    every stage; norm defaults to RMSNorm.
+
+    Papers: PVT/PVTv2 (Wang et al., 2021/2022, https://arxiv.org/abs/2102.12122,
+    https://arxiv.org/abs/2106.13797) for the SRA pyramid; Swin (Liu et al., 2021,
+    https://arxiv.org/abs/2103.14030) for the window-attention stages.
+    """
+    n = len(dims)
+    stages: list[PyramidStageConfig] = []
+    for i in range(n):
+        dim, dim_head = dims[i], dims[i] // heads[i]
+        attn: SpatialAttnConfig = (
+            WindowAttention2DConfig(
+                dim=dim, heads=heads[i], dim_head=dim_head, window=window, dropout=dropout
+            )
+            if i < window_stages
+            else SpatialReductionAttentionConfig(
+                dim=dim, heads=heads[i], dim_head=dim_head, reduction=sr_ratios[i], dropout=dropout
+            )
+        )
+        stages.append(
+            PyramidStageConfig(
+                attn=attn,
+                ff=SwiGLUConfig(dim=dim, mult=ff_mult, dropout=dropout),
+                norm=RMSNormConfig(dim=dim),
+                pos_encoding=RoPE2DConfig(dim_head=dim_head),
+                depth=depths[i],
+            )
+        )
+    merges = [PatchMergingConfig(in_dim=dims[i], out_dim=dims[i + 1]) for i in range(n - 1)]
+    return PyramidVisionConfig(stages=stages, inter_stage=merges)
+
+
+class _PyramidStage(nn.Module):
+    """A stack of SpatialTransformerLayers at a fixed grid; SpatialInput → SpatialInput."""
+
+    def __init__(self, config: PyramidStageConfig) -> None:
+        super().__init__()
+        pos = build_pos_encoding(config.pos_encoding)
+        self.layers = nn.ModuleList(
+            SpatialTransformerLayer(
+                attn=build_spatial_attn(config.attn, pos),
+                ff=build_ff(config.ff),
+                norm_attn=build_norm(config.norm),
+                norm_ff=build_norm(config.norm),
+            )
+            for _ in range(config.depth)
+        )
+
+    def forward(self, input: SpatialInput) -> SpatialInput:
+        for layer in self.layers:
+            input = layer(input)
+        return input
+
+
+def _to_feature_map(input: SpatialInput) -> Float[Tensor, "b d h w"]:
+    h, w = input.grid
+    return rearrange(input.x, "b (h w) d -> b d h w", h=h, w=w)
+
+
+class PyramidVisionBackbone(nn.Module):
+    """Multi-scale vision backbone (PVTv2-style hybrid).
+
+    Consumes a SpatialInput from your own patch-embed stem and returns one feature map per
+    stage as (b, d, H, W) — the multi-scale pyramid a detection/segmentation head expects.
+
+    Paper: "PVTv2: Improved Baselines with Pyramid Vision Transformer", Wang et al., 2022 —
+    https://arxiv.org/abs/2106.13797
+    """
+
+    def __init__(self, config: PyramidVisionConfig) -> None:
+        super().__init__()
+        self.stages = nn.ModuleList(_PyramidStage(stage) for stage in config.stages)
+        self.merges = nn.ModuleList(build_patch_merging(m) for m in config.inter_stage)
+
+    def forward(self, input: SpatialInput) -> list[Float[Tensor, "b d h w"]]:
+        """Shape flow (4-stage example, stem grid H₀×W₀, dims d₀..d₃):
+
+        stem in   (b, H₀·W₀, d₀)
+        stage i   (b, Hᵢ·Wᵢ, dᵢ)        emit feature map (b, dᵢ, Hᵢ, Wᵢ)
+        merge i   (b, Hᵢ₊₁·Wᵢ₊₁, dᵢ₊₁)  grid halved, width widened
+        returns   [(b, dᵢ, Hᵢ, Wᵢ) for each stage]  — the feature pyramid
+        """
+        features: list[Tensor] = []
+        x = input
+        for i, stage in enumerate(self.stages):
+            x = stage(x)
+            features.append(_to_feature_map(x))
+            if i < len(self.merges):
+                x = self.merges[i](x)
+        return features

stackformers-3.9.0/stackformers/spatial/README.md

@@ -0,0 +1,49 @@
+# spatial
+
+Attention over a 2-D grid of image tokens, for efficient vision backbones. Self-contained:
+nothing here touches the 1-D `SequenceInput` / `attention/` stack — it only reuses
+`padded_sdpa`, position encodings, norms, and feed-forwards.
+
+## Design
+
+**A separate input type, not a widened `SequenceInput`.** `SpatialInput` carries the grid
+shape `(H, W)` so attention can reshape the flat token sequence back to `(b, d, H, W)`. It is
+deliberately kept out of the `SequenceInput` union: that union stays a sealed 1-D abstraction,
+and spatial attention gets its own `SpatialAttn` protocol. Padded-only — fixed-resolution
+images do not vary in length, so there is no packed spatial variant.
+
+**Two attention variants, one protocol.** Both satisfy `SpatialAttn` and stay dot-product
+attention (so RoPE 2-D, GQA, and `qk_norm` all apply):
+- `WindowAttention2D` — non-overlapping `w×w` tiles, O(n·w²). Local; cheapest at the
+  high-resolution early stages. Distinct from `SelfAttentionConfig.window_size`, which is a
+  *1-D sliding* window over a flat sequence; this partitions the *2-D grid* into square tiles.
+- `SpatialReductionAttention` — full-grid queries attend to a spatially-reduced K/V context
+  (PVTv2 SRA), O(n·s) with s ≪ n. Global; affordable once the grid is small.
+
+**Reduction is an injected `KVReduction`.** `SpatialReductionAttention` delegates K/V
+downsampling to a collaborator (`NoKVReduction` for r=1 full attention, `ConvKVReduction` for a
+strided conv). The reduction's norm is itself an injected `Norm` — pluggable like everywhere
+else, defaulting to RMSNorm.
+
+**Window inverse is a closure.** `partition_windows` returns `(windows, merge)`; the `merge`
+inverse is bound to the exact `(b, H, W, window)` used to partition, so the un-partition step
+can never be called with mismatched parameters.
+
+**Export-friendly.** Grid shapes are Python ints throughout (positions, window counts, merged
+grid sizes), so reshapes stay static under `torch.export` / `torch.compile`.
+
+## Extending
+
+Add a new variant as a new `nn.Module` satisfying `SpatialAttn`, plus a config with a `kind`
+discriminator and a `match` arm in `factory.py`. To add a new K/V reduction, implement
+`KVReduction` and extend `build_kv_reduction`. Shifted windows would be a new variant here, not
+a flag on `WindowAttention2D`.
+
+The opinionated multi-scale assembly (stages + patch-merging → feature pyramid) lives in
+`presets/pyramid_vision.py`, not here — this module only provides the reusable primitives.
+
+## References
+
+- **Swin Transformer** — window partitioning. Liu et al., 2021. https://arxiv.org/abs/2103.14030
+- **Pyramid Vision Transformer (PVT)** — spatial-reduction attention. Wang et al., 2021. https://arxiv.org/abs/2102.12122
+- **PVTv2** — overlapping patch embedding / improved pyramid. Wang et al., 2022. https://arxiv.org/abs/2106.13797

stackformers-3.9.0/stackformers/spatial/__init__.py

@@ -0,0 +1,45 @@
+from stackformers.spatial.config import (
+    ConvKVReductionConfig,
+    KVReductionConfig,
+    NoKVReductionConfig,
+    PatchMergingConfig,
+    SpatialAttnConfig,
+    SpatialReductionAttentionConfig,
+    WindowAttention2DConfig,
+)
+from stackformers.spatial.factory import (
+    build_kv_reduction,
+    build_patch_merging,
+    build_spatial_attn,
+)
+from stackformers.spatial.input import SpatialInput, grid_positions, make_spatial_input
+from stackformers.spatial.kv_reduction import ConvKVReduction, NoKVReduction
+from stackformers.spatial.layer import SpatialTransformerLayer
+from stackformers.spatial.patch_merging import PatchMerging
+from stackformers.spatial.protocols import KVReduction, SpatialAttn
+from stackformers.spatial.spatial_reduction import SpatialReductionAttention
+from stackformers.spatial.window import WindowAttention2D
+
+__all__ = [
+    "SpatialInput",
+    "make_spatial_input",
+    "grid_positions",
+    "SpatialAttn",
+    "KVReduction",
+    "NoKVReduction",
+    "ConvKVReduction",
+    "WindowAttention2D",
+    "SpatialReductionAttention",
+    "SpatialTransformerLayer",
+    "PatchMerging",
+    "NoKVReductionConfig",
+    "ConvKVReductionConfig",
+    "KVReductionConfig",
+    "WindowAttention2DConfig",
+    "SpatialReductionAttentionConfig",
+    "SpatialAttnConfig",
+    "PatchMergingConfig",
+    "build_kv_reduction",
+    "build_spatial_attn",
+    "build_patch_merging",
+]

stackformers-3.9.0/stackformers/spatial/config.py

@@ -0,0 +1,138 @@
+from __future__ import annotations
+
+from typing import Annotated, Literal
+
+from pydantic import BaseModel, Field, model_validator
+
+from stackformers.attention.config import _validate_attn_dims
+from stackformers.norm.config import NormConfig, RMSNormConfig
+
+
+class NoKVReductionConfig(BaseModel):
+    kind: Literal["none"] = "none"
+
+
+class ConvKVReductionConfig(BaseModel):
+    """Strided-conv spatial reduction of the key/value context.
+
+    Paper: "Pyramid Vision Transformer" (SRA), Wang et al., 2021 —
+    https://arxiv.org/abs/2102.12122
+    """
+
+    kind: Literal["conv"] = "conv"
+    dim: int = Field(gt=0, description="Channel width of the tokens being reduced.")
+    reduction: int = Field(
+        gt=1,
+        description="Spatial stride r: an H×W grid is pooled to (H/r)×(W/r). Must divide H and W.",
+    )
+    norm: NormConfig = Field(
+        default_factory=lambda: RMSNormConfig(dim=1),
+        description="Norm applied to the reduced tokens. Defaults to RMSNorm, matching the "
+        "rest of the library; set to LayerNormConfig for the original PVTv2 behaviour.",
+    )
+
+    @model_validator(mode="after")
+    def _sync_norm_dim(self) -> ConvKVReductionConfig:
+        # The norm always acts on `dim` channels, so derive its width rather than making
+        # the caller repeat it (and risk a mismatch).
+        self.norm = self.norm.model_copy(update={"dim": self.dim})
+        return self
+
+
+KVReductionConfig = Annotated[
+    NoKVReductionConfig | ConvKVReductionConfig,
+    Field(discriminator="kind"),
+]
+
+
+class _SpatialAttnBase(BaseModel):
+    """Shared geometry for 2-D attention variants (mirrors SelfAttentionConfig)."""
+
+    dim: int = Field(gt=0, description="Model dimension — input/output width of the sublayer.")
+    heads: int = Field(default=8, gt=0, description="Number of query heads.")
+    dim_head: int = Field(default=64, gt=0, description="Dimension per head.")
+    kv_heads: int | None = Field(
+        default=None, description="Key/value heads. None = heads (MHA); a divisor of heads = GQA."
+    )
+    dropout: float = Field(default=0.0, ge=0.0, le=1.0, description="Output-projection dropout.")
+    qk_norm: bool = Field(default=False, description="RMSNorm queries and keys before the product.")
+
+    @model_validator(mode="after")
+    def _validate_dims(self) -> _SpatialAttnBase:
+        _validate_attn_dims(self.dim, self.heads, self.dim_head, self.kv_heads)
+        return self
+
+    @property
+    def effective_kv_heads(self) -> int:
+        return self.kv_heads if self.kv_heads is not None else self.heads
+
+    @property
+    def groups(self) -> int:
+        return self.heads // self.effective_kv_heads
+
+
+class WindowAttention2DConfig(_SpatialAttnBase):
+    """Non-overlapping w×w window attention over the grid — O(n·w²), local.
+
+    Distinct from SelfAttentionConfig.window_size, which is a 1-D sliding window over a
+    flat sequence; this partitions the 2-D grid into square tiles.
+
+    Paper: "Swin Transformer" (window partitioning), Liu et al., 2021 —
+    https://arxiv.org/abs/2103.14030
+    """
+
+    kind: Literal["window2d"] = "window2d"
+    window: int = Field(gt=0, description="Window side length w; must divide grid H and W.")
+
+
+class SpatialReductionAttentionConfig(_SpatialAttnBase):
+    """Global attention with spatially-reduced K/V — O(n·s).
+
+    Paper: "Pyramid Vision Transformer" (SRA), Wang et al., 2021 —
+    https://arxiv.org/abs/2102.12122
+    """
+
+    kind: Literal["spatial_reduction"] = "spatial_reduction"
+    reduction: int = Field(
+        default=1,
+        ge=1,
+        description="K/V spatial stride r. 1 = full global attention (no reduction).",
+    )
+    kv_norm: NormConfig = Field(
+        default_factory=lambda: RMSNormConfig(dim=1),
+        description="Norm applied to reduced K/V tokens (ignored when reduction == 1).",
+    )
+
+    @model_validator(mode="after")
+    def _sync_kv_norm_dim(self) -> SpatialReductionAttentionConfig:
+        self.kv_norm = self.kv_norm.model_copy(update={"dim": self.dim})
+        return self
+
+
+SpatialAttnConfig = Annotated[
+    WindowAttention2DConfig | SpatialReductionAttentionConfig,
+    Field(discriminator="kind"),
+]
+
+
+class PatchMergingConfig(BaseModel):
+    """Overlapping strided-conv downsample between pyramid stages (defaults k3 s2 p1).
+
+    Paper: "PVTv2: Improved Baselines with Pyramid Vision Transformer" (overlapping patch
+    embedding), Wang et al., 2022 — https://arxiv.org/abs/2106.13797
+    """
+
+    in_dim: int = Field(gt=0, description="Input channel width.")
+    out_dim: int = Field(gt=0, description="Output channel width after merging.")
+    kernel: int = Field(default=3, gt=0, description="Conv kernel size.")
+    stride: int = Field(default=2, gt=0, description="Conv stride — the downsample factor.")
+    padding: int = Field(default=1, ge=0, description="Conv padding.")
+    norm: NormConfig = Field(
+        default_factory=lambda: RMSNormConfig(dim=1),
+        description="Norm applied to merged tokens; dim is synced to out_dim.",
+    )
+
+    @model_validator(mode="after")
+    def _sync_norm_dim(self) -> PatchMergingConfig:
+        self.norm = self.norm.model_copy(update={"dim": self.out_dim})
+        return self

stackformers-3.9.0/stackformers/spatial/factory.py

@@ -0,0 +1,49 @@
+from __future__ import annotations
+
+from stackformers.norm.factory import build_norm
+from stackformers.positional.protocols import PosEncoding
+from stackformers.spatial.config import (
+    ConvKVReductionConfig,
+    KVReductionConfig,
+    NoKVReductionConfig,
+    PatchMergingConfig,
+    SpatialAttnConfig,
+    SpatialReductionAttentionConfig,
+    WindowAttention2DConfig,
+)
+from stackformers.spatial.kv_reduction import ConvKVReduction, NoKVReduction
+from stackformers.spatial.patch_merging import PatchMerging
+from stackformers.spatial.protocols import KVReduction, SpatialAttn
+from stackformers.spatial.spatial_reduction import SpatialReductionAttention
+from stackformers.spatial.window import WindowAttention2D
+
+
+def build_kv_reduction(config: KVReductionConfig) -> KVReduction:
+    match config:
+        case NoKVReductionConfig():
+            return NoKVReduction()
+        case ConvKVReductionConfig():
+            return ConvKVReduction(config, build_norm(config.norm))
+        case _:
+            raise AssertionError(f"Unhandled KV reduction config: {type(config)}")
+
+
+def _kv_reduction_for(config: SpatialReductionAttentionConfig) -> KVReductionConfig:
+    if config.reduction == 1:
+        return NoKVReductionConfig()
+    return ConvKVReductionConfig(dim=config.dim, reduction=config.reduction, norm=config.kv_norm)
+
+
+def build_spatial_attn(config: SpatialAttnConfig, pos_encoding: PosEncoding) -> SpatialAttn:
+    match config:
+        case WindowAttention2DConfig():
+            return WindowAttention2D(config, pos_encoding)
+        case SpatialReductionAttentionConfig():
+            reducer = build_kv_reduction(_kv_reduction_for(config))
+            return SpatialReductionAttention(config, pos_encoding, reducer)
+        case _:
+            raise AssertionError(f"Unhandled spatial attention config: {type(config)}")
+
+
+def build_patch_merging(config: PatchMergingConfig) -> PatchMerging:
+    return PatchMerging(config, build_norm(config.norm))

stackformers-3.9.0/stackformers/spatial/input.py

@@ -0,0 +1,53 @@
+from __future__ import annotations
+
+from typing import NamedTuple
+
+import torch
+from jaxtyping import Bool, Float
+from torch import Tensor
+
+
+def grid_positions(h: int, w: int, device: torch.device) -> Float[Tensor, "n c"]:
+    """Row-major (row, col) coordinates for an h×w grid, shape (h*w, 2).
+
+    Built with static h, w so it stays torch.export / torch.compile friendly.
+    """
+    rows = torch.arange(h, device=device).repeat_interleave(w)
+    cols = torch.arange(w, device=device).repeat(h)
+    return torch.stack([rows, cols], dim=-1).to(torch.float32)
+
+
+class SpatialInput(NamedTuple):
+    """Padded batch of tokens laid out on a 2-D grid.
+
+    Distinct from SequenceInput: it carries the grid shape so spatial attention can
+    reshape the flat token sequence back to (b, d, H, W). Padded-only — there is no
+    packed spatial variant (fixed-resolution images do not vary in length).
+
+    NamedTuple (not dataclass) so torch.export / torch.compile see the tensor fields
+    natively, matching PaddedInput / PackedInput.
+    """
+
+    x: Float[Tensor, "b n d"]
+    mask: Bool[Tensor, "b n"]  # True = valid token
+    abs_positions: Float[Tensor, "b n c"]  # c=2, (row, col) per token
+    grid: tuple[int, int]  # (H, W); H * W == n
+
+
+def make_spatial_input(
+    x: Float[Tensor, "b n d"],
+    grid: tuple[int, int],
+    mask: Bool[Tensor, "b n"] | None = None,
+) -> SpatialInput:
+    """Build a SpatialInput with row-major (row, col) positions from the grid shape.
+
+    mask defaults to all-valid; pass an explicit mask only when the batch is padded.
+    """
+    h, w = grid
+    n = h * w
+    if x.shape[1] != n:
+        raise ValueError(f"x has {x.shape[1]} tokens but grid {grid} implies {n}")
+    positions = grid_positions(h, w, x.device).unsqueeze(0).expand(x.shape[0], -1, -1)
+    if mask is None:
+        mask = torch.ones(x.shape[0], n, dtype=torch.bool, device=x.device)
+    return SpatialInput(x=x, mask=mask, abs_positions=positions, grid=grid)

stackformers-3.9.0/stackformers/spatial/kv_reduction.py

@@ -0,0 +1,69 @@
+from __future__ import annotations
+
+import torch.nn as nn
+from einops import rearrange
+from jaxtyping import Float
+from torch import Tensor
+
+from stackformers.norm.protocols import Norm
+from stackformers.spatial.config import ConvKVReductionConfig
+from stackformers.spatial.input import grid_positions
+
+
+class NoKVReduction(nn.Module):
+    """Null object for KVReduction — returns the full grid unchanged.
+
+    Keeps the attention forward() free of `if reduction is not None` branches.
+    """
+
+    def __call__(
+        self,
+        x: Float[Tensor, "b n d"],
+        grid: tuple[int, int],
+    ) -> tuple[Float[Tensor, "b s d"], Float[Tensor, "b s c"]]:
+        h, w = grid
+        positions = grid_positions(h, w, x.device).unsqueeze(0).expand(x.shape[0], -1, -1)
+        return x, positions
+
+
+class ConvKVReduction(nn.Module):
+    """Downsample K/V context with a strided conv, then norm.
+
+    Reshapes the flat token grid to (b, d, H, W), applies a stride-r conv that pools each
+    r×r block to one token, flattens back, and normalises with the injected Norm. Emits the
+    coarse grid's (row, col) positions so the caller can position-encode the reduced keys.
+
+    Paper: "Pyramid Vision Transformer" (spatial-reduction attention), Wang et al., 2021 —
+    https://arxiv.org/abs/2102.12122
+    """
+
+    def __init__(self, config: ConvKVReductionConfig, norm: Norm) -> None:
+        super().__init__()
+        self.reduction = config.reduction
+        self.conv = nn.Conv2d(
+            config.dim, config.dim, kernel_size=config.reduction, stride=config.reduction
+        )
+        self.norm = norm
+
+    def __call__(
+        self,
+        x: Float[Tensor, "b n d"],
+        grid: tuple[int, int],
+    ) -> tuple[Float[Tensor, "b s d"], Float[Tensor, "b s c"]]:
+        """Shape flow (r = reduction, s = (H/r)·(W/r)):
+
+        x        (b, n, d)      n = H·W
+        to grid  (b, d, H, W)
+        conv     (b, d, H/r, W/r)
+        ctx      (b, s, d)      flattened + normed reduced tokens
+        pos      (b, s, 2)      coarse-grid (row, col)
+        """
+        h, w = grid
+        r = self.reduction
+        if h % r != 0 or w % r != 0:
+            raise ValueError(f"grid {grid} not divisible by reduction {r}")
+        feat = rearrange(x, "b (h w) d -> b d h w", h=h, w=w)
+        feat = self.conv(feat)
+        ctx = self.norm(rearrange(feat, "b d h w -> b (h w) d"))
+        positions = grid_positions(h // r, w // r, x.device).unsqueeze(0).expand(x.shape[0], -1, -1)
+        return ctx, positions

stackformers-3.9.0/stackformers/spatial/layer.py

@@ -0,0 +1,42 @@
+from __future__ import annotations
+
+import torch.nn as nn
+
+from stackformers.feedforward.protocols import FeedForward
+from stackformers.norm.protocols import Norm
+from stackformers.spatial.input import SpatialInput
+from stackformers.spatial.protocols import SpatialAttn
+
+
+class SpatialTransformerLayer(nn.Module):
+    """Pre-norm transformer layer over a 2-D grid: norm → spatial-attn → residual,
+    norm → ff → residual.
+
+    Mirrors TransformerLayer but is typed to SpatialAttn / SpatialInput so the grid shape
+    rides through unchanged via NamedTuple._replace. The feed-forward and norms are the same
+    injected protocols used everywhere else.
+    """
+
+    def __init__(
+        self,
+        attn: SpatialAttn,
+        ff: FeedForward,
+        norm_attn: Norm,
+        norm_ff: Norm,
+    ) -> None:
+        super().__init__()
+        self.attn = attn
+        self.ff = ff
+        self.norm_attn = norm_attn
+        self.norm_ff = norm_ff
+
+    def forward(self, input: SpatialInput) -> SpatialInput:
+        """Shape-invariant: x stays (b, n, d) throughout; grid rides along via _replace.
+
+        x = x + attn(norm(x))   # (b, n, d) → (b, n, d)
+        x = x + ff(norm(x))     # (b, n, d) → (b, n, d)
+        """
+        normed = input._replace(x=self.norm_attn(input.x))
+        x = input.x + self.attn(normed)
+        x = x + self.ff(self.norm_ff(x))
+        return input._replace(x=x)