stackformers 3.8.2__tar.gz → 3.9.1__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: stackformers
-Version: 3.8.2
+Version: 3.9.1
 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.1}/pyproject.toml

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

{stackformers-3.8.2 → stackformers-3.9.1}/stackformers/attention/README.md

@@ -8,9 +8,11 @@ Multi-head self-attention and cross-attention.
 
 **Static vs runtime dispatch.** The attention *type* (global vs sliding-window) is determined at construction time via `window_size: int | None` in `SelfAttentionConfig`. The *backend* (padded vs packed) is determined at runtime: `SelfAttention.forward` and `CrossAttention.forward` dispatch on `SequenceInput` via `match`. This means one model handles both training (packed) and inference (padded) without swapping classes.
 
-**Kernel ops in `ops.py`.** All low-level attention math lives in `attention/ops.py`: mask builders, `packed_attn` (thin `varlen_attn` wrapper), `padded_sdpa`, and `packed_attn_or_fallback`. The modules stay thin dispatchers.
+**Kernel ops in `ops.py`.** All low-level attention math lives in `attention/ops.py`: mask builders, `padded_sdpa`, the packed↔padded scatter/gather helpers, and `packed_attn_or_fallback`. The modules stay thin dispatchers.
 
-**Packed attention with automatic fallback.** `PackedInput` normally routes to `varlen_attn`, which requires CUDA and `float16`/`bfloat16`. When those conditions are not met (CPU, float32, `torch.export`), `packed_attn_or_fallback` transparently scatters q/k/v to padded layout, runs `F.scaled_dot_product_attention`, and gathers the result back to packed — same weights, same output shape `(nt, h, d)`. A `UserWarning` is emitted so callers know they are on the slow path. Dropout works on the fallback path; it is silently skipped only when `varlen_attn` is used (which does not accept a dropout argument).
+**Experimental kernel behind one interface function.** `torch.nn.attention.varlen` is an experimental API — the symbol is absent in some torch builds and its signature is not stable. Everything that touches it lives in `varlen_backend.py` behind a single function, `try_varlen_attn(...) -> Tensor | None`: the guarded import, the eligibility/bias checks, and the call itself. The rest of the package depends only on that function and never on the volatile import path, so an absent or renamed symbol degrades to the fallback instead of breaking the whole package at import time.
+
+**Packed attention with automatic fallback.** `PackedInput` normally routes to `varlen_attn`, which requires CUDA and `float16`/`bfloat16`. `packed_attn_or_fallback` calls `try_varlen_attn`; a `None` result means take the padded path — scatter q/k/v to padded layout, run `F.scaled_dot_product_attention`, and gather the result back to packed (same weights, same output shape `(nt, h, d)`). The fallback runs **silently** for the expected CPU/float32/`torch.export` route, and **with a `UserWarning` explaining why** when the kernel was eligible but unusable: it failed to import, an attention bias is present (no bias slot), or the call raised a runtime/signature error (unsupported GPU or a changed experimental API). Dropout works on the fallback path; it is silently skipped only when `varlen_attn` is used (which does not accept a dropout argument).
 
 **`torch.export` and `torch.compile` compatibility.** All scatter/gather helpers in `ops.py` (`_cu_to_indices`, `_packed_heads_to_padded`, `_padded_heads_to_packed`) are fully export-compatible. The key design constraint is that per-document length iteration must never use Python-side loops over data-dependent values (i.e., no `.tolist()` followed by a list comprehension). Instead, `_cu_to_indices` uses `torch.repeat_interleave` to produce `batch_idx` and derives `pos_idx` via the identity `arange(nt) - cu[batch_idx]`. This keeps the computation graph purely symbolic even when `nt` (total tokens) is a data-dependent unbacked symbol.
 

{stackformers-3.8.2 → stackformers-3.9.1}/stackformers/attention/ops.py

@@ -1,19 +1,13 @@
 from __future__ import annotations
 
-import warnings
-
 import torch
 import torch.nn.functional as F
 from torch import Tensor
-from torch.nn.attention.varlen import varlen_attn as _varlen_attn
 
+from stackformers.attention.varlen_backend import try_varlen_attn
 from stackformers.sequence import PackedSequence
 
 
-def varlen_supported(q: Tensor) -> bool:
-    return q.is_cuda and q.dtype in (torch.float16, torch.bfloat16)
-
-
 def padding_mask(mask: Tensor, dtype: torch.dtype) -> Tensor:
     bias = torch.zeros(mask.shape, dtype=dtype, device=mask.device)
     bias.masked_fill_(~mask, torch.finfo(dtype).min)
@@ -34,12 +28,6 @@ def window_mask(n: int, s: int, window_size: int, causal: bool, device: torch.de
     return mask
 
 
-def varlen_window(causal: bool, window_size: int | None) -> tuple[int, int]:
-    if window_size is None:
-        return (-1, 0) if causal else (-1, -1)
-    return (window_size, 0) if causal else (window_size // 2, window_size // 2)
-
-
 def padded_sdpa(
     q: Tensor,
     k: Tensor,
@@ -101,30 +89,6 @@ def _padded_heads_to_packed(x: Tensor, mask: Tensor) -> Tensor:
     return x.permute(0, 2, 1, 3)[mask]
 
 
-def packed_attn(
-    q: Tensor,
-    k: Tensor,
-    v: Tensor,
-    q_seq: PackedSequence,
-    k_seq: PackedSequence,
-    causal: bool,
-    window_size: int | None,
-) -> Tensor:
-    win = varlen_window(causal, window_size)
-    result = _varlen_attn(
-        query=q,
-        key=k,
-        value=v,
-        cu_seq_q=q_seq.cu_seqlens.to(torch.int32),
-        cu_seq_k=k_seq.cu_seqlens.to(torch.int32),
-        max_q=q_seq.max_seqlen,
-        max_k=k_seq.max_seqlen,
-        window_size=win,
-    )
-    assert isinstance(result, Tensor)
-    return result
-
-
 def packed_attn_or_fallback(
     q: Tensor,
     k: Tensor,
@@ -135,29 +99,16 @@ def packed_attn_or_fallback(
     window_size: int | None,
     bias: Tensor | None,
 ) -> Tensor:
-    """Run varlen attention; fall back to padded SDPA when unavailable.
+    """Run varlen attention; fall back to padded SDPA when it is unavailable.
 
     Inputs and output are all packed: q/k/v are (nt, h, d), result is (nt_q, h, d).
-    bias=None → no attention bias; varlen_attn is tried first.
-    bias=Tensor → attention bias present; varlen_attn is skipped (it has no bias slot) and a
-    warning is emitted. Fallback also activates on CPU, non-float16/bfloat16 dtypes, or GPUs
-    that do not support varlen_attn (e.g. compute capability < 8.0).
+    ``try_varlen_attn`` owns the experimental kernel and its warnings; ``None`` from it means
+    take the padded SDPA path — silently for the expected CPU/float32/export route, or after
+    a warning when the kernel was eligible but unavailable.
     """
-    if bias is not None and varlen_supported(q):
-        warnings.warn(
-            "varlen_attn does not support attention bias; falling back to padded SDPA. "
-            "Performance may be lower.",
-            stacklevel=2,
-        )
-    elif varlen_supported(q):
-        try:
-            return packed_attn(q, k, v, q_seq, k_seq, causal, window_size)
-        except RuntimeError as exc:
-            warnings.warn(
-                f"varlen_attn is not supported on this device ({exc}); "
-                "falling back to padded SDPA. Performance may be lower.",
-                stacklevel=2,
-            )
+    out = try_varlen_attn(q, k, v, q_seq, k_seq, causal, window_size, bias)
+    if out is not None:
+        return out
 
     # Keep b as SymInt-friendly: int(SymInt) forces specialization under torch.export.
     # Downstream consumers (_packed_heads_to_padded) accept SymInt for shape args.

stackformers-3.9.1/stackformers/attention/varlen_backend.py

@@ -0,0 +1,88 @@
+"""The experimental ``varlen_attn`` kernel behind one interface function.
+
+``torch.nn.attention.varlen`` is an experimental API: the symbol is absent in some torch
+builds, and its call signature is not yet stable. Everything that touches it lives here —
+the guarded import, the eligibility/bias checks, and the call itself — so the rest of the
+package depends only on :func:`try_varlen_attn` and never on the volatile import path.
+"""
+
+from __future__ import annotations
+
+import warnings
+
+import torch
+from torch import Tensor
+
+from stackformers.sequence import PackedSequence
+
+_FALLBACK_NOTE = " Falling back to padded SDPA; performance may be lower."
+
+
+def _load():
+    try:
+        from torch.nn.attention.varlen import varlen_attn
+    except ImportError as exc:
+        return None, f"varlen_attn could not be imported ({exc})"
+    return varlen_attn, None
+
+
+_varlen_attn, _import_error = _load()
+
+
+def _eligible(q: Tensor) -> bool:
+    """Whether this query's device/dtype is one ``varlen_attn`` would handle at all."""
+    return q.is_cuda and q.dtype in (torch.float16, torch.bfloat16)
+
+
+def _window(causal: bool, window_size: int | None) -> tuple[int, int]:
+    if window_size is None:
+        return (-1, 0) if causal else (-1, -1)
+    return (window_size, 0) if causal else (window_size // 2, window_size // 2)
+
+
+def _skip(reason: str) -> None:
+    warnings.warn(reason + _FALLBACK_NOTE, stacklevel=3)
+    return None
+
+
+def try_varlen_attn(
+    q: Tensor,
+    k: Tensor,
+    v: Tensor,
+    q_seq: PackedSequence,
+    k_seq: PackedSequence,
+    causal: bool,
+    window_size: int | None,
+    bias: Tensor | None,
+) -> Tensor | None:
+    """Attempt varlen attention for a packed batch; return ``None`` to use the SDPA fallback.
+
+    q/k/v are ``(nt, h, d)``; the result is the packed ``(nt_q, h, d)`` tensor. ``None`` means
+    the caller should run the padded SDPA path instead. When varlen *would* apply (CUDA +
+    float16/bfloat16) but cannot, a ``UserWarning`` explains why: an absent kernel, an
+    attention bias (varlen_attn has no bias slot), or a runtime/signature error from the call
+    (an unsupported GPU or a changed experimental API). The expected CPU/float32/export route
+    returns ``None`` silently — that is normal, not a failure.
+    """
+    if not _eligible(q):
+        return None
+    if _varlen_attn is None:
+        return _skip(f"varlen_attn is unavailable ({_import_error}).")
+    if bias is not None:
+        return _skip("varlen_attn does not support attention bias.")
+    try:
+        result = _varlen_attn(
+            query=q,
+            key=k,
+            value=v,
+            cu_seq_q=q_seq.cu_seqlens.to(torch.int32),
+            cu_seq_k=k_seq.cu_seqlens.to(torch.int32),
+            max_q=q_seq.max_seqlen,
+            max_k=k_seq.max_seqlen,
+            window_size=_window(causal, window_size),
+        )
+    except (RuntimeError, TypeError) as exc:
+        return _skip(f"varlen_attn call failed ({type(exc).__name__}: {exc}).")
+    if not isinstance(result, Tensor):
+        return _skip(f"varlen_attn returned {type(result).__name__}, expected Tensor.")
+    return result

{stackformers-3.8.2 → stackformers-3.9.1}/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.1/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.1/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.1/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.1/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