stackformers 4.2.0__tar.gz → 4.2.1__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (87) hide show
  1. {stackformers-4.2.0 → stackformers-4.2.1}/CHANGELOG.md +34 -0
  2. {stackformers-4.2.0 → stackformers-4.2.1}/PKG-INFO +1 -1
  3. {stackformers-4.2.0 → stackformers-4.2.1}/pyproject.toml +1 -1
  4. {stackformers-4.2.0 → stackformers-4.2.1}/stackformers/positional/README.md +6 -4
  5. {stackformers-4.2.0 → stackformers-4.2.1}/stackformers/positional/config.py +20 -18
  6. {stackformers-4.2.0 → stackformers-4.2.1}/stackformers/positional/rope_nd.py +18 -8
  7. {stackformers-4.2.0 → stackformers-4.2.1}/tests/positional/test_rope_nd.py +35 -9
  8. {stackformers-4.2.0 → stackformers-4.2.1}/uv.lock +1 -1
  9. {stackformers-4.2.0 → stackformers-4.2.1}/.claudeignore +0 -0
  10. {stackformers-4.2.0 → stackformers-4.2.1}/.gitignore +0 -0
  11. {stackformers-4.2.0 → stackformers-4.2.1}/.python-version +0 -0
  12. {stackformers-4.2.0 → stackformers-4.2.1}/.vscode/settings.json +0 -0
  13. {stackformers-4.2.0 → stackformers-4.2.1}/CLAUDE.md +0 -0
  14. {stackformers-4.2.0 → stackformers-4.2.1}/Justfile +0 -0
  15. {stackformers-4.2.0 → stackformers-4.2.1}/LICENSE +0 -0
  16. {stackformers-4.2.0 → stackformers-4.2.1}/PLAN.md +0 -0
  17. {stackformers-4.2.0 → stackformers-4.2.1}/README.md +0 -0
  18. {stackformers-4.2.0 → stackformers-4.2.1}/run-remote-tests.sh +0 -0
  19. {stackformers-4.2.0 → stackformers-4.2.1}/stackformers/__init__.py +0 -0
  20. {stackformers-4.2.0 → stackformers-4.2.1}/stackformers/attention/README.md +0 -0
  21. {stackformers-4.2.0 → stackformers-4.2.1}/stackformers/attention/__init__.py +0 -0
  22. {stackformers-4.2.0 → stackformers-4.2.1}/stackformers/attention/bias.py +0 -0
  23. {stackformers-4.2.0 → stackformers-4.2.1}/stackformers/attention/config.py +0 -0
  24. {stackformers-4.2.0 → stackformers-4.2.1}/stackformers/attention/cross_attn.py +0 -0
  25. {stackformers-4.2.0 → stackformers-4.2.1}/stackformers/attention/distance_bias.py +0 -0
  26. {stackformers-4.2.0 → stackformers-4.2.1}/stackformers/attention/factory.py +0 -0
  27. {stackformers-4.2.0 → stackformers-4.2.1}/stackformers/attention/ops.py +0 -0
  28. {stackformers-4.2.0 → stackformers-4.2.1}/stackformers/attention/protocols.py +0 -0
  29. {stackformers-4.2.0 → stackformers-4.2.1}/stackformers/attention/self_attn.py +0 -0
  30. {stackformers-4.2.0 → stackformers-4.2.1}/stackformers/attention/varlen_backend.py +0 -0
  31. {stackformers-4.2.0 → stackformers-4.2.1}/stackformers/config.py +0 -0
  32. {stackformers-4.2.0 → stackformers-4.2.1}/stackformers/cross_attender.py +0 -0
  33. {stackformers-4.2.0 → stackformers-4.2.1}/stackformers/decoder.py +0 -0
  34. {stackformers-4.2.0 → stackformers-4.2.1}/stackformers/encoder.py +0 -0
  35. {stackformers-4.2.0 → stackformers-4.2.1}/stackformers/feedforward/README.md +0 -0
  36. {stackformers-4.2.0 → stackformers-4.2.1}/stackformers/feedforward/__init__.py +0 -0
  37. {stackformers-4.2.0 → stackformers-4.2.1}/stackformers/feedforward/config.py +0 -0
  38. {stackformers-4.2.0 → stackformers-4.2.1}/stackformers/feedforward/factory.py +0 -0
  39. {stackformers-4.2.0 → stackformers-4.2.1}/stackformers/feedforward/geglu.py +0 -0
  40. {stackformers-4.2.0 → stackformers-4.2.1}/stackformers/feedforward/protocols.py +0 -0
  41. {stackformers-4.2.0 → stackformers-4.2.1}/stackformers/feedforward/relu_squared.py +0 -0
  42. {stackformers-4.2.0 → stackformers-4.2.1}/stackformers/feedforward/swiglu.py +0 -0
  43. {stackformers-4.2.0 → stackformers-4.2.1}/stackformers/layers.py +0 -0
  44. {stackformers-4.2.0 → stackformers-4.2.1}/stackformers/norm/README.md +0 -0
  45. {stackformers-4.2.0 → stackformers-4.2.1}/stackformers/norm/__init__.py +0 -0
  46. {stackformers-4.2.0 → stackformers-4.2.1}/stackformers/norm/config.py +0 -0
  47. {stackformers-4.2.0 → stackformers-4.2.1}/stackformers/norm/factory.py +0 -0
  48. {stackformers-4.2.0 → stackformers-4.2.1}/stackformers/norm/protocols.py +0 -0
  49. {stackformers-4.2.0 → stackformers-4.2.1}/stackformers/positional/__init__.py +0 -0
  50. {stackformers-4.2.0 → stackformers-4.2.1}/stackformers/positional/factory.py +0 -0
  51. {stackformers-4.2.0 → stackformers-4.2.1}/stackformers/positional/learned.py +0 -0
  52. {stackformers-4.2.0 → stackformers-4.2.1}/stackformers/positional/none.py +0 -0
  53. {stackformers-4.2.0 → stackformers-4.2.1}/stackformers/positional/protocols.py +0 -0
  54. {stackformers-4.2.0 → stackformers-4.2.1}/stackformers/positional/rope1d.py +0 -0
  55. {stackformers-4.2.0 → stackformers-4.2.1}/stackformers/positional/rope2d.py +0 -0
  56. {stackformers-4.2.0 → stackformers-4.2.1}/stackformers/presets/README.md +0 -0
  57. {stackformers-4.2.0 → stackformers-4.2.1}/stackformers/presets/__init__.py +0 -0
  58. {stackformers-4.2.0 → stackformers-4.2.1}/stackformers/presets/cross_attender.py +0 -0
  59. {stackformers-4.2.0 → stackformers-4.2.1}/stackformers/presets/decoder.py +0 -0
  60. {stackformers-4.2.0 → stackformers-4.2.1}/stackformers/presets/encoder.py +0 -0
  61. {stackformers-4.2.0 → stackformers-4.2.1}/stackformers/sequence.py +0 -0
  62. {stackformers-4.2.0 → stackformers-4.2.1}/tests/__init__.py +0 -0
  63. {stackformers-4.2.0 → stackformers-4.2.1}/tests/attention/__init__.py +0 -0
  64. {stackformers-4.2.0 → stackformers-4.2.1}/tests/attention/test_cross_attn.py +0 -0
  65. {stackformers-4.2.0 → stackformers-4.2.1}/tests/attention/test_distance_bias.py +0 -0
  66. {stackformers-4.2.0 → stackformers-4.2.1}/tests/attention/test_kernels.py +0 -0
  67. {stackformers-4.2.0 → stackformers-4.2.1}/tests/attention/test_ops.py +0 -0
  68. {stackformers-4.2.0 → stackformers-4.2.1}/tests/attention/test_self_attn.py +0 -0
  69. {stackformers-4.2.0 → stackformers-4.2.1}/tests/attention/test_varlen_backend.py +0 -0
  70. {stackformers-4.2.0 → stackformers-4.2.1}/tests/conftest.py +0 -0
  71. {stackformers-4.2.0 → stackformers-4.2.1}/tests/feedforward/__init__.py +0 -0
  72. {stackformers-4.2.0 → stackformers-4.2.1}/tests/feedforward/test_geglu.py +0 -0
  73. {stackformers-4.2.0 → stackformers-4.2.1}/tests/feedforward/test_relu_squared.py +0 -0
  74. {stackformers-4.2.0 → stackformers-4.2.1}/tests/feedforward/test_swiglu.py +0 -0
  75. {stackformers-4.2.0 → stackformers-4.2.1}/tests/positional/__init__.py +0 -0
  76. {stackformers-4.2.0 → stackformers-4.2.1}/tests/positional/test_learned_pos.py +0 -0
  77. {stackformers-4.2.0 → stackformers-4.2.1}/tests/positional/test_rope.py +0 -0
  78. {stackformers-4.2.0 → stackformers-4.2.1}/tests/presets/__init__.py +0 -0
  79. {stackformers-4.2.0 → stackformers-4.2.1}/tests/presets/test_cross_attender.py +0 -0
  80. {stackformers-4.2.0 → stackformers-4.2.1}/tests/presets/test_decoder.py +0 -0
  81. {stackformers-4.2.0 → stackformers-4.2.1}/tests/presets/test_encoder.py +0 -0
  82. {stackformers-4.2.0 → stackformers-4.2.1}/tests/test_cross_attender.py +0 -0
  83. {stackformers-4.2.0 → stackformers-4.2.1}/tests/test_decoder.py +0 -0
  84. {stackformers-4.2.0 → stackformers-4.2.1}/tests/test_encoder.py +0 -0
  85. {stackformers-4.2.0 → stackformers-4.2.1}/tests/test_layers.py +0 -0
  86. {stackformers-4.2.0 → stackformers-4.2.1}/tests/test_norm.py +0 -0
  87. {stackformers-4.2.0 → stackformers-4.2.1}/tests/test_sequence.py +0 -0
@@ -6,6 +6,40 @@ The format follows [Keep a Changelog](https://keepachangelog.com/en/1.1.0/), and
6
6
  adheres to [Semantic Versioning](https://semver.org/): MAJOR for breaking public API changes,
7
7
  MINOR for backwards-compatible features, PATCH for bug fixes and internal changes.
8
8
 
9
+ ## [4.2.1] — 2026-07-14
10
+
11
+ ### Removed
12
+
13
+ - **`RoPENDConfig.headroom`.** It was redundant with `r_max`: the ladder only ever depended on
14
+ the product `headroom · r_max`, so the two knobs shared one degree of freedom, and anyone
15
+ wanting more reach at the slow end could get it by raising `r_max` — the parameter that has a
16
+ measurement procedure attached. It was also the last *tuned* number in a module whose premise
17
+ is that the band range is measured, not tuned.
18
+
19
+ Its default has a derivation, so it need not be a parameter at all. Attention sees *signed*
20
+ offsets spanning `[−r_max, +r_max]`, a width of `2·r_max`; asking that the slowest band turn
21
+ through at most half a circle across that width gives `ω_lo = π / (2·r_max)` directly. That is
22
+ bit-identical to the old `headroom=4.0` default at every `r_max`, so callers on the default —
23
+ which is to say every caller, since the field shipped in 4.2.0 — see no change in behaviour.
24
+
25
+ Both ends of the ladder are now the same rule, `ω = π / scale`: a half turn over `r_min` at
26
+ the fast end (Nyquist), and a half turn over `2·r_max` at the slow end (no wrap).
27
+
28
+ Treated as a PATCH rather than a MAJOR bump because `headroom` existed for exactly one
29
+ release, is not known to be set anywhere, and its removal cannot change the behaviour of a
30
+ caller that did not set it. Note that pydantic's default `extra="ignore"` means a leftover
31
+ `RoPENDConfig(headroom=...)` is *dropped silently* rather than raising — a caller who had set
32
+ it to something other than 4.0 will now get different frequencies without being told.
33
+
34
+ ### Fixed
35
+
36
+ - **`RoPENDConfig` now requires `dim_head >= 4 * coords`** (at least two bands per axis).
37
+ `dim_head == 2 * coords` divides cleanly and passed validation, but leaves a one-band ladder
38
+ with nowhere to descend: `torch.linspace(hi, lo, 1)` returns `[hi]`, so the lone band lands on
39
+ the fast end, `r_max` is discarded entirely, and the encoding becomes periodic with period
40
+ `2 · r_min` across the whole domain — offsets of `0`, `2·r_min`, `4·r_min` … all produce an
41
+ identical attention logit. It failed silently; it is now rejected at construction.
42
+
9
43
  ## [4.2.0] — 2026-07-14
10
44
 
11
45
  ### Added
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: stackformers
3
- Version: 4.2.0
3
+ Version: 4.2.1
4
4
  Summary: Typed, composable, SOLID transformer library for PyTorch
5
5
  Project-URL: Homepage, https://github.com/Red-Eyed/stackformers
6
6
  Project-URL: Repository, https://github.com/Red-Eyed/stackformers
@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
4
4
 
5
5
  [project]
6
6
  name = "stackformers"
7
- version = "4.2.0"
7
+ version = "4.2.1"
8
8
  description = "Typed, composable, SOLID transformer library for PyTorch"
9
9
  readme = "README.md"
10
10
  license = { file = "LICENSE" }
@@ -8,12 +8,14 @@ Positional encodings applied to Q and K tensors inside attention — not to the
8
8
 
9
9
  RoPE's shortest wavelength is fixed at 2π coordinate units, and `base` only stretches the long end of the band range. So `base` lands the ladder correctly only when tokens sit *one unit apart* — true for text, true for a patch grid, meaningless for scattered continuous coordinates. Positions normalised to `[0, 1]` leave every band longer than the whole domain, so the encoding barely rotates and carries almost no signal; and `base=10000` on a 14×14 patch grid leaves 14 of 16 bands frozen.
10
10
 
11
- `RotaryEmbeddingND` (`RoPENDConfig`) drops `base` entirely and builds the ladder from the two numbers that actually mean something, in any number of dimensions:
11
+ `RotaryEmbeddingND` (`RoPENDConfig`) drops `base` entirely and builds the ladder from the two numbers that actually mean something, in any number of dimensions. Both ends are the same rule — **a half turn over the scale in question, `ω = π / scale`**:
12
12
 
13
- - `r_min` — the finest separation the model must resolve. Sets the shortest wavelength to `2·r_min`, the Nyquist limit. Measure it as a low percentile of the nearest-neighbour distance.
14
- - `r_max` — the domain diameter. Sets the longest wavelength (times `headroom`, so the slowest band stays monotone rather than wrapping). Measure it as a high percentile of the pairwise distance distribution.
13
+ - `r_min` — the finest separation the model must resolve. Turned through π on the fastest band, so the shortest wavelength is `2·r_min`: the Nyquist limit, past which distinct offsets alias onto the same rotation. Measure it as a low percentile of the nearest-neighbour distance.
14
+ - `r_max` — the domain diameter. Attention sees *signed* offsets, spanning `[−r_max, +r_max]` — a width of `2·r_max`, not `r_max` — so the slowest band is given a half turn over that full width. It cannot wrap, which leaves it monotone in the offset: a coarse ramp the model reads as displacement. Measure it as a high percentile of the pairwise distance distribution.
15
15
 
16
- The resulting ladder depends only on the *dynamic range* `r_max / r_min`, so metres, pixels and millimetres all give the identical encoding — the property `base` never had. `RoPE1DConfig` and `RoPE2DConfig` keep `base` and are unchanged; use them for text and grids.
16
+ There is no third knob, and no constant to tune: two measured distances fix the ladder. The resulting band range depends only on the *dynamic range* `r_max / r_min`, so metres, pixels and millimetres all give the identical encoding — the property `base` never had. `RoPE1DConfig` and `RoPE2DConfig` keep `base` and are unchanged; use them for text and grids.
17
+
18
+ `dim_head` must be at least `4 · coords`, so that every axis gets at least two bands. At one band per axis the ladder has nowhere to descend: it collapses onto the fast end, `r_max` is never reached, and the encoding silently becomes periodic with period `2·r_min` across the whole domain. `RoPENDConfig` rejects that at construction.
17
19
 
18
20
  **Centre your coordinates.** Not for invariance — RoPE is translation-invariant exactly, since the query and key rotations cancel into `ω·(pᵢ − pⱼ)`. It is for float32: the angle is `ω·p`, and coordinates far from the origin push it into the hundreds of radians and spend the mantissa before the cosine is taken. A 1e5 offset costs three orders of magnitude of accuracy. This cannot be done inside the module — cross-attention must subtract the *same* constant from query and key positions, and the module sees them separately — so subtract the centroid at the input boundary.
19
21
 
@@ -53,9 +53,10 @@ class RoPENDConfig(BaseModel):
53
53
  Takes no ``base``. RoPE's ``base`` ladder pins its fastest band at a wavelength of 2π
54
54
  *coordinate units* regardless of the base, so it only lands correctly when tokens sit one
55
55
  unit apart — true for text and patch grids, meaningless for continuous coordinates. The
56
- band range is set by the data instead: r_min fixes the fast end, r_max the slow end. The
57
- ladder then depends only on the ratio r_max / r_min, so the units the coordinates happen
58
- to be expressed in stop mattering.
56
+ band range is set by the data instead: r_min fixes the fast end, r_max the slow end, and
57
+ each is a half turn over the scale it names (ω = π / scale). Nothing else is free. The
58
+ ladder's shape then depends only on the ratio r_max / r_min, so the units the coordinates
59
+ happen to be expressed in stop mattering.
59
60
  """
60
61
 
61
62
  kind: Literal["rope_nd"] = "rope_nd"
@@ -66,29 +67,22 @@ class RoPENDConfig(BaseModel):
66
67
  r_min: float = Field(
67
68
  gt=0.0,
68
69
  description=(
69
- "Finest separation between two nodes that the model must tell apart. Sets the"
70
- " shortest wavelength to 2*r_min — the Nyquist limit, below which distinct offsets"
71
- " alias onto the same rotation. Measure it as a low percentile of the"
72
- " nearest-neighbour distance, not the minimum, which is noise."
70
+ "Finest separation between two nodes that the model must tell apart. Gets a half"
71
+ " turn on the fastest band, so the shortest wavelength is 2*r_min — the Nyquist"
72
+ " limit, below which distinct offsets alias onto the same rotation. Measure it as a"
73
+ " low percentile of the nearest-neighbour distance, not the minimum, which is noise."
73
74
  ),
74
75
  )
75
76
  r_max: float = Field(
76
77
  gt=0.0,
77
78
  description=(
78
- "Diameter of the domain — the largest offset that must stay distinguishable. Sets"
79
- " the longest wavelength. Measure it as a high percentile of the pairwise distance"
79
+ "Diameter of the domain — the largest offset that must stay distinguishable."
80
+ " Attention sees signed offsets spanning [-r_max, +r_max], so the slowest band is"
81
+ " given a half turn over that full width of 2*r_max: it never wraps, which leaves it"
82
+ " monotone in the offset. Measure it as a high percentile of the pairwise distance"
80
83
  " distribution, not the maximum, which is an outlier."
81
84
  ),
82
85
  )
83
- headroom: float = Field(
84
- default=4.0,
85
- gt=1.0,
86
- description=(
87
- "How far the longest wavelength reaches past r_max. Keeps the slowest band monotone"
88
- " across the whole domain instead of wrapping, so it can act as a coarse absolute"
89
- " coordinate. Llama's defaults sit at roughly 13x; 2-4x is ample here."
90
- ),
91
- )
92
86
 
93
87
  @model_validator(mode="after")
94
88
  def _check(self) -> RoPENDConfig:
@@ -99,6 +93,14 @@ class RoPENDConfig(BaseModel):
99
93
  f" dim_head is split into {self.coords} per-axis blocks, and each block is"
100
94
  " rotated in pairs of channels."
101
95
  )
96
+ if self.bands_per_axis < 2:
97
+ raise ValueError(
98
+ f"dim_head ({self.dim_head}) leaves {self.bands_per_axis} band per axis; at"
99
+ f" least 2 bands are needed, so dim_head must be at least {4 * self.coords}."
100
+ " A lone band lands on the fast end of the ladder and r_max is never reached,"
101
+ " leaving the encoding periodic with period 2 * r_min across the whole domain:"
102
+ " every offset an integer number of periods apart becomes indistinguishable."
103
+ )
102
104
  if self.r_max <= self.r_min:
103
105
  raise ValueError(f"r_max ({self.r_max}) must exceed r_min ({self.r_min}).")
104
106
  return self
@@ -18,18 +18,28 @@ def _frequency_ladder(config: RoPENDConfig) -> Tensor:
18
18
  at ω=1, i.e. a wavelength of exactly 2π coordinate units, whatever the base — so ``base``
19
19
  can only stretch the slow end, and the ladder lands correctly only when tokens happen to
20
20
  sit one unit apart. That holds for text and for patch grids; it is meaningless for
21
- continuous coordinates, where the useful band range is fixed by the data instead:
21
+ continuous coordinates, where the band range is fixed by the data instead.
22
22
 
23
- ω_hi = π / r_min → shortest wavelength = 2·r_min (Nyquist on the
24
- finest separation the model must resolve)
25
- ω_lo = 2π / (headroom · r_max) → longest wavelength spans the whole domain
23
+ Both ends are a half turn over the scale they name — ω = π / scale:
26
24
 
27
- So the ladder is a function of the *dynamic range* r_max / r_min alone, and is invariant
28
- to the units the coordinates are expressed in metres, pixels, or normalised to [0, 1]
29
- all give the same encoding. That is precisely the property ``base`` does not have.
25
+ ω_hi = π / r_min half-period r_min. The finest separation the model must resolve
26
+ is turned through π, the furthest two points can be driven apart
27
+ on a circle. Any faster is past Nyquist, and distinct offsets
28
+ start aliasing onto the same rotation.
29
+
30
+ ω_lo = π / (2·r_max) half-period 2·r_max. Attention sees *signed* offsets, spanning
31
+ [−r_max, +r_max] — a width of 2·r_max, not r_max — so the slowest
32
+ band turns through at most π across the entire domain. It cannot
33
+ wrap, which leaves it monotone in the offset: a coarse ramp the
34
+ model can read directly as displacement.
35
+
36
+ Nothing here is tuned, and there is no knob left to set: two measured distances fix the
37
+ ladder. Its *shape* depends only on the dynamic range r_max / r_min, so the encoding is
38
+ invariant to the units the coordinates are expressed in — metres, pixels, or normalised to
39
+ [0, 1] all give the same result. That is precisely the property ``base`` does not have.
30
40
  """
31
41
  omega_hi = math.pi / config.r_min
32
- omega_lo = 2.0 * math.pi / (config.headroom * config.r_max)
42
+ omega_lo = math.pi / (2.0 * config.r_max)
33
43
  steps = torch.linspace(math.log(omega_hi), math.log(omega_lo), config.bands_per_axis)
34
44
  return torch.exp(steps) # fastest band first, matching the 1-D/2-D convention
35
45
 
@@ -11,7 +11,7 @@ from stackformers.positional.rope_nd import RotaryEmbeddingND
11
11
  from tests.conftest import atol
12
12
 
13
13
  B, N, H = 2, 16, 4
14
- R_MIN, R_MAX, HEADROOM = 0.5, 100.0, 4.0
14
+ R_MIN, R_MAX = 0.5, 100.0
15
15
 
16
16
  DIMS = [(1, 64), (2, 64), (3, 96)] # (coords, dim_head) — dim_head must divide by 2 * coords
17
17
 
@@ -19,9 +19,7 @@ DIMS = [(1, 64), (2, 64), (3, 96)] # (coords, dim_head) — dim_head must divid
19
19
  @pytest.fixture(params=DIMS, ids=lambda p: f"c{p[0]}")
20
20
  def config(request: pytest.FixtureRequest) -> RoPENDConfig:
21
21
  coords, dim_head = request.param
22
- return RoPENDConfig(
23
- dim_head=dim_head, coords=coords, r_min=R_MIN, r_max=R_MAX, headroom=HEADROOM
24
- )
22
+ return RoPENDConfig(dim_head=dim_head, coords=coords, r_min=R_MIN, r_max=R_MAX)
25
23
 
26
24
 
27
25
  @pytest.fixture
@@ -65,15 +63,32 @@ def test_output_shape(
65
63
 
66
64
 
67
65
  def test_ladder_spans_nyquist_to_domain(config: RoPENDConfig) -> None:
68
- """The whole point of dropping `base`: the band range is pinned to the data, not a constant."""
66
+ """The whole point of dropping `base`: the band range is pinned to the data, not a constant.
67
+
68
+ Both ends are a half turn over the scale they name — r_min at the fast end, and the full
69
+ signed offset width 2*r_max at the slow end — so both are stated here as half-periods.
70
+ """
69
71
  inv_freq: torch.Tensor = RotaryEmbeddingND(config).inv_freq # type: ignore[assignment]
70
- shortest = 2 * math.pi / float(inv_freq.max())
71
- longest = 2 * math.pi / float(inv_freq.min())
72
- assert shortest == pytest.approx(2 * R_MIN, rel=1e-4) # Nyquist on the finest separation
73
- assert longest == pytest.approx(HEADROOM * R_MAX, rel=1e-4) # spans the domain, with headroom
72
+ fastest_half_period = math.pi / float(inv_freq.max())
73
+ slowest_half_period = math.pi / float(inv_freq.min())
74
+ assert fastest_half_period == pytest.approx(R_MIN, rel=1e-4) # Nyquist on the finest gap
75
+ assert slowest_half_period == pytest.approx(2 * R_MAX, rel=1e-4) # the whole signed range
74
76
  assert inv_freq.shape == (config.bands_per_axis,)
75
77
 
76
78
 
79
+ def test_slowest_band_never_wraps(config: RoPENDConfig) -> None:
80
+ """The slow end must stay monotone over every offset attention can see.
81
+
82
+ Attention sees signed offsets over [-r_max, +r_max], so the slowest band sweeps
83
+ omega_lo * 2 * r_max. Let that exceed pi and the band turns back on itself: two different
84
+ offsets in the domain collide on the same rotation, in the one band whose entire job is to
85
+ tell the coarse end apart.
86
+ """
87
+ inv_freq: torch.Tensor = RotaryEmbeddingND(config).inv_freq # type: ignore[assignment]
88
+ sweep = float(inv_freq.min()) * 2 * R_MAX
89
+ assert sweep == pytest.approx(math.pi, rel=1e-4)
90
+
91
+
77
92
  def test_units_do_not_matter(config: RoPENDConfig, device: torch.device) -> None:
78
93
  """Scale the coordinates and r_min/r_max together — the encoding must not move.
79
94
 
@@ -209,6 +224,17 @@ def test_rejects_indivisible_head_dim() -> None:
209
224
  RoPENDConfig(dim_head=64, coords=3, r_min=R_MIN, r_max=R_MAX)
210
225
 
211
226
 
227
+ def test_rejects_single_band_per_axis() -> None:
228
+ """dim_head == 2 * coords divides cleanly, but leaves a ladder with nowhere to descend.
229
+
230
+ The lone band lands on omega_hi and r_max is never reached, so the encoding is periodic with
231
+ period 2 * r_min across the entire domain — offsets of 0, 2*r_min, 4*r_min ... all give the
232
+ same rotation. It fails silently, which is why it is rejected at construction.
233
+ """
234
+ with pytest.raises(ValidationError, match="at least 2 bands"):
235
+ RoPENDConfig(dim_head=6, coords=3, r_min=R_MIN, r_max=R_MAX)
236
+
237
+
212
238
  def test_rejects_inverted_range() -> None:
213
239
  with pytest.raises(ValidationError, match="must exceed r_min"):
214
240
  RoPENDConfig(dim_head=64, coords=2, r_min=100.0, r_max=1.0)
@@ -811,7 +811,7 @@ wheels = [
811
811
 
812
812
  [[package]]
813
813
  name = "stackformers"
814
- version = "4.2.0"
814
+ version = "4.2.1"
815
815
  source = { editable = "." }
816
816
  dependencies = [
817
817
  { name = "beartype" },
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes