stackformers 3.9.1__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.1 → stackformers-3.9.2}/PKG-INFO

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

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

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

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