stackformers 3.9.0__tar.gz → 3.9.2__tar.gz

stackformers-3.9.2/CHANGELOG.md

@@ -0,0 +1,219 @@
+# Changelog
+
+All notable changes to this project are documented here.
+
+The format follows [Keep a Changelog](https://keepachangelog.com/en/1.1.0/), and this project
+adheres to [Semantic Versioning](https://semver.org/): MAJOR for breaking public API changes,
+MINOR for backwards-compatible features, PATCH for bug fixes and internal changes.
+
+## [3.9.2] — 2026-07-14
+
+### Fixed
+
+- **`RotaryEmbedding2D` was not computing a rotation.** The frequency vector was laid out as
+  `[row | row | col | col]`, but `_rotate_half` pairs channel `i` with channel `i + dh/2`, so a
+  row angle was paired against a column angle. The resulting 2×2 map had
+  `det = cos((row − col) · ωᵢ)` and singular values `√(1 ± |sin D|)` — a squeeze, not a rotation.
+  It was therefore non-orthogonal, annihilated high-frequency channels wherever the determinant
+  reached zero, and produced attention scores that leaked absolute position instead of depending
+  only on relative offset. Frequencies are now built as `[row | col | row | col]`.
+- `RotaryEmbedding2D` now casts positions and `inv_freq` to float32 for the outer product,
+  matching `RotaryEmbedding1D`. A `.half()`-ed module previously lost position precision.
+
+### Changed
+
+- **Numerics of `RotaryEmbedding2D` change.** Any checkpoint trained against the previous
+  encoding is tuned to the broken position code and requires retraining or finetuning.
+
+### Added
+
+- Math tests for `RotaryEmbedding2D` on a genuine row-major grid: a definition-first reference
+  oracle, orthogonality (norms and inner products), relative-offset invariance, per-axis
+  sensitivity, distinct-offset separation, and packed/padded agreement. The previous test used a
+  `row == col` position grid — the main diagonal, the one locus where the row and column angles
+  coincide and the broken layout is indistinguishable from a correct one.
+
+## [3.9.1] — 2026-06-30
+
+### Fixed
+
+- Make the experimental `varlen_attn` import and call robust across PyTorch versions.
+
+## [3.9.0] — 2026-06-17
+
+### Added
+
+- 2-D grid attention (`spatial`) and the pyramid vision backbone preset.
+
+## [3.8.2] — 2026-05-29
+
+### Fixed
+
+- Force absolute positions to float32 in `sequence`.
+
+## [3.8.1] — 2026-05-29
+
+### Fixed
+
+- Keep batch size as a `SymInt` in the packed attention fallback, preserving `torch.export`
+  compatibility.
+
+## [3.8.0] — 2026-05-27
+
+### Fixed
+
+- Replace `.tolist()` loops in `sequence` and `ops` with tensor operations, making them
+  `torch.export`-compatible.
+
+### Added
+
+- PyPI publish metadata in `pyproject.toml`.
+
+## [3.7.0] — 2026-05-21
+
+### Changed
+
+- Introduce abstract generic base classes for all presets.
+
+## [3.6.0] — 2026-05-21
+
+### Added
+
+- `AttnBias` protocol and the `NoAttnBias` null object on `SelfAttention`.
+
+## [3.5.2] — 2026-05-21
+
+### Fixed
+
+- Warn and fall back to SDPA when `varlen_attn` raises.
+
+## [3.5.1] — 2026-05-18
+
+### Changed
+
+- Move attention dropout to the output projection.
+
+## [3.5.0] — 2026-05-18
+
+### Added
+
+- Packed→padded SDPA fallback; extract `ops.py`.
+
+## [3.4.2] — 2026-05-18
+
+### Added
+
+- `padded_to_packed` and `packed_to_padded` helpers.
+
+## [3.4.1] — 2026-05-18
+
+### Removed
+
+- Packed-attention CPU fallback; warn on dropout instead.
+
+## [3.4.0] — 2026-05-17
+
+### Added
+
+- ReLU² feed-forward network.
+
+## [3.3.0] — 2026-05-17
+
+### Added
+
+- QK-norm in attention; all config fields documented.
+
+## [3.2.0] — 2026-05-17
+
+### Changed
+
+- Eliminate layout-specific duplicates in positional encoding implementations.
+
+## [3.1.0] — 2026-05-17
+
+### Changed
+
+- Replace `forward()` on positional encodings with `forward_padded` / `forward_packed`, so
+  callers pick the path for the layout they already know.
+
+## [3.0.1] — 2026-05-17
+
+### Fixed
+
+- Use the GELU tanh approximation in GEGLU.
+
+## [3.0.0] — 2026-05-16
+
+### Changed
+
+- **Breaking:** remove the kernel abstraction; padded and packed dispatch are unified.
+
+## [2.0.2] — 2026-05-16
+
+### Removed
+
+- The unused local-attention package.
+
+## [2.0.1] — 2026-05-16
+
+### Changed
+
+- Store feed-forward activations as `nn.Module` attributes.
+
+## [2.0.0] — 2026-05-16
+
+### Removed
+
+- **Breaking:** the `attn_bias` / ALiBi system.
+
+## [1.1.0] — 2026-05-16
+
+### Added
+
+- `LearnedPosEncoding` with absolute position embeddings.
+
+## [1.0.1] — 2026-05-16
+
+### Added
+
+- `LayerNorm` wrapper.
+
+## [1.0.0] — 2026-05-16
+
+### Added
+
+- GEGLU feed-forward variant.
+
+### Changed
+
+- **Breaking:** feed-forward config becomes a discriminated union.
+
+## [0.5.0] — 2026-05-16
+
+### Added
+
+- Unfold mode for the windowed SDPA kernel.
+
+## [0.4.0] — 2026-05-15
+
+### Added
+
+- `RotaryEmbedding2D` and the packed cross-attender.
+
+## [0.3.0] — 2026-05-14
+
+### Changed
+
+- Establish the semantic versioning convention.
+
+## [0.2.0] — 2026-05-13
+
+### Changed
+
+- Remove the `v1` namespace; add the `presets` package.
+
+## [0.1.0] — 2026-05-13
+
+### Added
+
+- Initial library scaffold.

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

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

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

{stackformers-3.9.0 → stackformers-3.9.2}/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.9.0 → stackformers-3.9.2}/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.2/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.9.0 → stackformers-3.9.2}/stackformers/positional/rope2d.py

@@ -25,12 +25,22 @@ class RotaryEmbedding2D(nn.Module):
 
     @torch.no_grad()
     def _freqs_from_positions(self, positions: Tensor) -> Tensor:
-        """positions: (..., 2) → freqs: (..., dh)"""
+        """positions: (..., 2) → freqs: (..., dh), laid out as [row | col | row | col].
+
+        _rotate_half pairs channel i with channel i + dh/2, so both members of a pair must
+        carry the same angle for the map to be a rotation. Duplicating the *concatenated*
+        half-vector satisfies that; concatenating the *duplicated* row and col blocks
+        ([row | row | col | col]) would pair a row angle against a col angle and yield a
+        squeeze instead — non-orthogonal, and not a function of relative position.
+
+        float32 cast ensures half-precision inputs don't lose precision in the outer product.
+        """
         inv: Tensor = self.inv_freq  # type: ignore[assignment]
-        row = positions[..., 0].to(dtype=inv.dtype).unsqueeze(-1) * inv  # (..., dh//4)
-        col = positions[..., 1].to(dtype=inv.dtype).unsqueeze(-1) * inv
-        half = lambda f: torch.cat([f, f], dim=-1)  # noqa: E731
-        return torch.cat([half(row), half(col)], dim=-1)  # (..., dh)
+        pos = positions.to(dtype=torch.float32)
+        row = pos[..., 0].unsqueeze(-1) * inv.float()  # (..., dh//4)
+        col = pos[..., 1].unsqueeze(-1) * inv.float()  # (..., dh//4)
+        half = torch.cat([row, col], dim=-1)  # (..., dh//2)
+        return torch.cat([half, half], dim=-1)  # (..., dh)
 
     def _encode(
         self, q: Tensor, k: Tensor, q_positions: Tensor, k_positions: Tensor

stackformers-3.9.2/tests/attention/test_varlen_backend.py

@@ -0,0 +1,122 @@
+"""Tests for stackformers/attention/varlen_backend.py.
+
+Covers the fallback contract of ``try_varlen_attn``: it returns the kernel's tensor on the
+happy path and ``None`` (with a reason-bearing ``UserWarning`` when the kernel was eligible
+but unusable) on every degraded path — absent symbol, attention bias, a runtime/signature
+error from the call, or a non-Tensor return. The experimental kernel is faked via monkeypatch
+so the logic is exercised without CUDA or the real (possibly-absent) ``varlen_attn``.
+"""
+
+from __future__ import annotations
+
+import warnings
+
+import pytest
+import torch
+
+from stackformers.attention import varlen_backend as vb
+from stackformers.sequence import PackedSequence
+
+NT, H, DH = 5, 2, 8
+
+Packed = tuple[torch.Tensor, torch.Tensor, torch.Tensor, PackedSequence]
+
+
+@pytest.fixture
+def packed() -> Packed:
+    """Two packed sequences (lengths 3 and 2) of (nt, h, d) q/k/v plus their PackedSequence."""
+    seq = PackedSequence(cu_seqlens=torch.tensor([0, 3, 5], dtype=torch.long), max_seqlen=3)
+    q = torch.randn(NT, H, DH)
+    k = torch.randn(NT, H, DH)
+    v = torch.randn(NT, H, DH)
+    return q, k, v, seq
+
+
+@pytest.fixture
+def force_eligible(monkeypatch: pytest.MonkeyPatch) -> None:
+    """Pretend the query is CUDA + float16 so the eligibility guard does not short-circuit."""
+    monkeypatch.setattr(vb, "_eligible", lambda q: True)
+
+
+def _attempt(packed: Packed, bias: torch.Tensor | None = None) -> torch.Tensor | None:
+    q, k, v, seq = packed
+    return vb.try_varlen_attn(q, k, v, seq, seq, causal=False, window_size=None, bias=bias)
+
+
+def test_ineligible_falls_back_silently(packed: Packed, recwarn: pytest.WarningsRecorder) -> None:
+    """A plain CPU/float32 query is the expected fallback route — None, and no warning."""
+    out = _attempt(packed)
+    assert out is None
+    assert len(recwarn) == 0
+
+
+def test_unavailable_kernel_warns(
+    packed: Packed, force_eligible: None, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    """An eligible query with no imported kernel falls back and reports the import reason."""
+    monkeypatch.setattr(vb, "_varlen_attn", None)
+    monkeypatch.setattr(vb, "_import_error", "varlen_attn could not be imported (no module)")
+    with pytest.warns(UserWarning, match="unavailable"):
+        out = _attempt(packed)
+    assert out is None
+
+
+def test_attention_bias_warns(
+    packed: Packed, force_eligible: None, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    """varlen_attn has no bias slot: an eligible query with a bias falls back with a warning."""
+    monkeypatch.setattr(vb, "_varlen_attn", lambda **kw: torch.zeros(NT, H, DH))
+    bias = torch.zeros(1, 1, NT, NT)
+    with pytest.warns(UserWarning, match="attention bias"):
+        out = _attempt(packed, bias=bias)
+    assert out is None
+
+
+def test_signature_error_warns(
+    packed: Packed, force_eligible: None, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    """A changed experimental signature raises TypeError, which is caught and falls back."""
+
+    def renamed_kwarg(**kw: object) -> torch.Tensor:
+        raise TypeError("unexpected keyword argument 'cu_seq_q'")
+
+    monkeypatch.setattr(vb, "_varlen_attn", renamed_kwarg)
+    with pytest.warns(UserWarning, match="call failed"):
+        out = _attempt(packed)
+    assert out is None
+
+
+def test_runtime_error_warns(
+    packed: Packed, force_eligible: None, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    """An unsupported-GPU RuntimeError from the call is caught and falls back."""
+
+    def unsupported(**kw: object) -> torch.Tensor:
+        raise RuntimeError("varlen_attn requires compute capability >= 8.0")
+
+    monkeypatch.setattr(vb, "_varlen_attn", unsupported)
+    with pytest.warns(UserWarning, match="call failed"):
+        out = _attempt(packed)
+    assert out is None
+
+
+def test_non_tensor_return_warns(
+    packed: Packed, force_eligible: None, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    """A changed return contract (non-Tensor) falls back rather than propagating bad data."""
+    monkeypatch.setattr(vb, "_varlen_attn", lambda **kw: ("not", "a", "tensor"))
+    with pytest.warns(UserWarning, match="expected Tensor"):
+        out = _attempt(packed)
+    assert out is None
+
+
+def test_success_returns_kernel_tensor(
+    packed: Packed, force_eligible: None, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    """The happy path returns the kernel's tensor unchanged and emits no warning."""
+    expected = torch.randn(NT, H, DH)
+    monkeypatch.setattr(vb, "_varlen_attn", lambda **kw: expected)
+    with warnings.catch_warnings():
+        warnings.simplefilter("error")
+        out = _attempt(packed)
+    assert out is expected

{stackformers-3.9.0 → stackformers-3.9.2}/tests/positional/test_rope.py

@@ -125,21 +125,176 @@ def test_rope1d_preserves_norms(
     assert torch.allclose(k.norm(dim=-1), k_out.norm(dim=-1), atol=tol)
 
 
+# --- RoPE-2D ---
+#
+# A row-major GH x GW grid, never row == col. On the diagonal (row == col) the row and column
+# angles coincide, which makes a broken frequency layout indistinguishable from a correct one.
+
+GH, GW = 3, 5
+GN = GH * GW
+
+
+def _grid_positions(b: int, device: torch.device, dtype: torch.dtype) -> torch.Tensor:
+    """Row-major (row, col) positions over a GH x GW grid, as (b, GN, 2)."""
+    rows = torch.arange(GH, device=device, dtype=dtype).repeat_interleave(GW)
+    cols = torch.arange(GW, device=device, dtype=dtype).repeat(GH)
+    return torch.stack([rows, cols], dim=-1).unsqueeze(0).expand(b, -1, -1).clone()
+
+
+def _reference_rope2d(t: torch.Tensor, positions: torch.Tensor, base: int = 10_000) -> torch.Tensor:
+    """Axial 2-D RoPE written directly from the definition, independent of the implementation.
+
+    Rotates channel pair (i, i + dh/2) by row * w_i for the first dh/4 pairs and by col * w_i
+    for the rest. t: (b, h, n, dh), positions: (b, n, 2).
+    """
+    dh = t.shape[-1]
+    half, quarter = dh // 2, dh // 4
+    inv = 1.0 / (base ** (torch.arange(0, half, 2, dtype=torch.float32) / half))
+    inv = inv.to(device=t.device)
+
+    out = torch.zeros_like(t, dtype=torch.float32)
+    src = t.float()
+    for i in range(half):
+        axis = 0 if i < quarter else 1
+        freq = inv[i if i < quarter else i - quarter]
+        angle = positions[..., axis].float() * freq  # (b, n)
+        cos, sin = angle.cos().unsqueeze(1), angle.sin().unsqueeze(1)  # (b, 1, n)
+        x, y = src[..., i], src[..., i + half]
+        out[..., i] = x * cos - y * sin
+        out[..., i + half] = x * sin + y * cos
+    return out.to(t.dtype)
+
+
+@pytest.fixture
+def grid_qk(device_dtype: tuple[torch.device, torch.dtype]) -> tuple[torch.Tensor, torch.Tensor]:
+    device, dtype = device_dtype
+    return (
+        torch.randn(B, H, GN, DH, device=device, dtype=dtype),
+        torch.randn(B, H, GN, DH, device=device, dtype=dtype),
+    )
+
+
+@pytest.fixture
+def grid_pos(device_dtype: tuple[torch.device, torch.dtype]) -> torch.Tensor:
+    device, dtype = device_dtype
+    return _grid_positions(B, device, dtype)
+
+
 def test_rope2d_output_shape(
     rope2d: RotaryEmbedding2D,
-    device_dtype: tuple[torch.device, torch.dtype],
+    grid_qk: tuple[torch.Tensor, torch.Tensor],
+    grid_pos: torch.Tensor,
 ) -> None:
-    device, dtype = device_dtype
-    q = torch.randn(B, H, N, DH, device=device, dtype=dtype)
-    k = torch.randn(B, H, N, DH, device=device, dtype=dtype)
-    row_ids = torch.arange(N, dtype=dtype, device=device)
-    col_ids = torch.arange(N, dtype=dtype, device=device)
-    pos = torch.stack([row_ids, col_ids], dim=-1).unsqueeze(0).expand(B, -1, -1)  # b n 2
-    q_out, k_out = rope2d.forward_padded(q, k, pos, pos)
+    q, k = grid_qk
+    q_out, k_out = rope2d.forward_padded(q, k, grid_pos, grid_pos)
     assert q_out.shape == q.shape
     assert k_out.shape == k.shape
 
 
+def test_rope2d_matches_reference_rotation(
+    rope2d: RotaryEmbedding2D,
+    grid_qk: tuple[torch.Tensor, torch.Tensor],
+    grid_pos: torch.Tensor,
+    device_dtype: tuple[torch.device, torch.dtype],
+) -> None:
+    """Pins the exact rotation against a definition-first oracle, not just its properties."""
+    _, dtype = device_dtype
+    q, k = grid_qk
+    q_out, k_out = rope2d.forward_padded(q, k, grid_pos, grid_pos)
+    assert torch.allclose(q_out, _reference_rope2d(q, grid_pos), atol=atol(dtype))
+    assert torch.allclose(k_out, _reference_rope2d(k, grid_pos), atol=atol(dtype))
+
+
+def test_rope2d_is_orthogonal(
+    rope2d: RotaryEmbedding2D,
+    grid_qk: tuple[torch.Tensor, torch.Tensor],
+    grid_pos: torch.Tensor,
+    device_dtype: tuple[torch.device, torch.dtype],
+) -> None:
+    """A rotation preserves norms and inner products; a squeeze or reflection does not."""
+    _, dtype = device_dtype
+    q, k = grid_qk
+    q_out, k_out = rope2d.forward_padded(q, k, grid_pos, grid_pos)
+    tol = atol(dtype)
+    assert torch.allclose(q.norm(dim=-1), q_out.norm(dim=-1), atol=tol)
+    assert torch.allclose(k.norm(dim=-1), k_out.norm(dim=-1), atol=tol)
+    # Same position for q and k, so the shared rotation must cancel in the inner product.
+    before = (q.float() * k.float()).sum(dim=-1)
+    after = (q_out.float() * k_out.float()).sum(dim=-1)
+    assert torch.allclose(before, after, atol=tol * DH)
+
+
+@pytest.mark.parametrize("offset", [(0, 1), (1, 0), (2, 3), (-1, 2)], ids=str)
+def test_rope2d_score_depends_only_on_relative_offset(
+    device: torch.device,
+    offset: tuple[int, int],
+) -> None:
+    """The defining property: sliding a query/key pair across the grid must not move the score."""
+    rope = RotaryEmbedding2D(RoPE2DConfig(dim_head=DH)).to(device=device)
+    q = torch.randn(1, H, 1, DH, device=device)
+    k = torch.randn(1, H, 1, DH, device=device)
+    d_row, d_col = offset
+
+    def score(row: int, col: int) -> torch.Tensor:
+        q_pos = torch.tensor([[[float(row), float(col)]]], device=device)
+        k_pos = torch.tensor([[[float(row + d_row), float(col + d_col)]]], device=device)
+        q_out, k_out = rope.forward_padded(q, k, q_pos, k_pos)
+        return (q_out @ k_out.transpose(-1, -2))[0, :, 0, 0]
+
+    anchor = score(0, 0)
+    for row, col in [(2, 3), (5, 9), (11, 4), (7, 7)]:
+        assert torch.allclose(anchor, score(row, col), atol=atol(torch.float32) * 10)
+
+
+def test_rope2d_distinct_offsets_give_distinct_scores(device: torch.device) -> None:
+    """Guards the degenerate fix where every offset collapses to the same score."""
+    rope = RotaryEmbedding2D(RoPE2DConfig(dim_head=DH)).to(device=device)
+    q = torch.randn(1, H, 1, DH, device=device)
+    k = torch.randn(1, H, 1, DH, device=device)
+
+    def score(k_row: float, k_col: float) -> torch.Tensor:
+        q_pos = torch.zeros(1, 1, 2, device=device)
+        k_pos = torch.tensor([[[k_row, k_col]]], device=device)
+        q_out, k_out = rope.forward_padded(q, k, q_pos, k_pos)
+        return (q_out @ k_out.transpose(-1, -2))[0, :, 0, 0]
+
+    # A row-offset and the matching col-offset must not be conflated: the axes are separate.
+    assert not torch.allclose(score(2.0, 0.0), score(0.0, 2.0))
+    assert not torch.allclose(score(0.0, 0.0), score(1.0, 1.0))
+
+
+@pytest.mark.parametrize("axis", [0, 1], ids=["row", "col"])
+def test_rope2d_each_axis_changes_the_encoding(
+    rope2d: RotaryEmbedding2D,
+    device_dtype: tuple[torch.device, torch.dtype],
+    axis: int,
+) -> None:
+    """Moving along either axis alone must change the output — neither axis may be ignored."""
+    device, dtype = device_dtype
+    q = torch.randn(1, H, 1, DH, device=device, dtype=dtype)
+    base_pos = torch.full((1, 1, 2), 2.0, device=device, dtype=dtype)
+    moved_pos = base_pos.clone()
+    moved_pos[..., axis] += 3.0
+    base_out, _ = rope2d.forward_padded(q, q, base_pos, base_pos)
+    moved_out, _ = rope2d.forward_padded(q, q, moved_pos, moved_pos)
+    assert not torch.allclose(base_out, moved_out, atol=atol(dtype))
+
+
+def test_rope2d_packed_matches_padded(
+    rope2d: RotaryEmbedding2D,
+    device_dtype: tuple[torch.device, torch.dtype],
+) -> None:
+    """Both layouts are the same per-token operation and must agree token for token."""
+    device, dtype = device_dtype
+    q = torch.randn(1, H, GN, DH, device=device, dtype=dtype)
+    pos = _grid_positions(1, device, dtype)
+    padded, _ = rope2d.forward_padded(q, q, pos, pos)
+
+    packed_q = q[0].transpose(0, 1).contiguous()  # (nt, h, dh)
+    packed_out, _ = rope2d.forward_packed(packed_q, packed_q, pos[0], pos[0])
+    assert torch.allclose(packed_out, padded[0].transpose(0, 1), atol=atol(dtype))
+
+
 # --- YaRN scaling ---
 
 _YARN = YaRNConfig(scale=4.0, original_max_seq_len=512)

{stackformers-3.9.0 → stackformers-3.9.2}/uv.lock

@@ -811,7 +811,7 @@ wheels = [
 
 [[package]]
 name = "stackformers"
-version = "3.9.0"
+version = "3.9.2"
 source = { editable = "." }
 dependencies = [
     { name = "beartype" },