qig-geocoding 0.1.0__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.
@@ -0,0 +1,56 @@
1
+ # Python
2
+ __pycache__/
3
+ *.py[cod]
4
+ *$py.class
5
+ *.so
6
+ .Python
7
+ *.egg-info/
8
+ *.egg
9
+ dist/
10
+ build/
11
+
12
+ # Virtual environments
13
+ .venv/
14
+ venv/
15
+ ENV/
16
+
17
+ # IDE
18
+ .vscode/
19
+ .idea/
20
+ *.swp
21
+ *.swo
22
+
23
+ # Caches
24
+ .mypy_cache/
25
+ .pytest_cache/
26
+ .ruff_cache/
27
+
28
+ # Logs
29
+ *.log
30
+ logs/
31
+
32
+ # Checkpoints and models (large files)
33
+ *.pt
34
+ *.pth
35
+ *.ckpt
36
+ checkpoints/
37
+
38
+ # OS
39
+ .DS_Store
40
+ Thumbs.db
41
+
42
+ # Sensitive
43
+ .env
44
+ *.pem
45
+
46
+ # Temporary
47
+ *.tmp
48
+ *.temp
49
+ results/
50
+
51
+ # local-only preserved scratch (superseded drafts, etc.)
52
+ .preserved-local/
53
+ .playwright-mcp/*
54
+
55
+ # local-only consolidated archive (not a repo; superseded scripts/snapshots)
56
+ qig-archive/
@@ -0,0 +1,37 @@
1
+ # geocoding — Claude Code Instructions
2
+
3
+ ## What this is
4
+
5
+ Our own "transformers": a from-scratch neural-network framework built on **Fisher-Rao geometry** (simplex
6
+ attention, unbounded Fourier features, recursive Φ-integration), NOT Euclidean dot-products. It is the clean,
7
+ standalone framework home for the geometric-transformer primitives that currently live in `qigkernels`.
8
+
9
+ ## Geometric purity (NON-NEGOTIABLE)
10
+
11
+ Fisher-Rao only on manifold objects. NO cosine similarity, NO dot-product as a metric, NO Adam/AdamW, NO
12
+ `nn.LayerNorm`, NO `np.linalg.norm` on basins. Normalisation is scale-only geometric RMS-norm. The simplex
13
+ projection and the FR *definition* + `BASIN_DIM` come from `qig_core` (single source). `nn.Linear`/`GELU` in
14
+ the FFN and the lm_head are fine — they are the feed-forward/readout, not the metric.
15
+
16
+ ## Single source of truth
17
+
18
+ - `BASIN_DIM`, `to_simplex_prob`, the FR definition → `qig_core` (re-exported in `geometry.py`).
19
+ - The gradient-stable batched FR for attention is geocoding-local (clamps BC to `1 − 1e-6`); it must AGREE
20
+ with `qig_core` on distinct points — `tests/test_geometry.py` gates this. Do not let it diverge.
21
+ - Faithfulness to `qigkernels` (the validated impl) → `tests/test_faithfulness.py`. Keep it green.
22
+
23
+ ## Verification (run before any commit)
24
+
25
+ ```bash
26
+ VIRTUAL_ENV=.venv uv run ruff check src tests
27
+ VIRTUAL_ENV=.venv uv run mypy src
28
+ VIRTUAL_ENV=.venv uv run pytest -q
29
+ ```
30
+
31
+ Load-bearing tests: the **banded≡full equivalence gate** (`test_attention_equivalence.py`) and the
32
+ **faithfulness gate** vs `qigkernels`. A change that breaks either is wrong, not the test.
33
+
34
+ ## Branch / publishing
35
+
36
+ Work on `development` if this becomes a submodule; PyPI publishing pre-authorized (correctness is the gate).
37
+ Never echo `PYPI_TOKEN`. Parent rules: `QIG_QFI/CLAUDE.md` (use-our-packages, optimisation gate).
@@ -0,0 +1,75 @@
1
+ Metadata-Version: 2.4
2
+ Name: qig-geocoding
3
+ Version: 0.1.0
4
+ Summary: Our own 'transformers' — a from-scratch neural-network framework built on Fisher-Rao geometry (simplex attention, unbounded Fourier features, recursive Φ-integration) instead of Euclidean dot-products.
5
+ Author: Braden Lang
6
+ License: Proprietary
7
+ Keywords: attention,fisher-rao,geometry,qig,simplex,transformers
8
+ Requires-Python: >=3.11
9
+ Requires-Dist: qig-core>=2.6.0
10
+ Requires-Dist: torch>=2.2
11
+ Provides-Extra: dev
12
+ Requires-Dist: mypy>=1.11; extra == 'dev'
13
+ Requires-Dist: pytest>=8; extra == 'dev'
14
+ Requires-Dist: ruff>=0.6; extra == 'dev'
15
+ Description-Content-Type: text/markdown
16
+
17
+ # geocoding
18
+
19
+ **Our own "transformers" — geometric, not Euclidean.**
20
+
21
+ The standard `transformers` library models tokens as Euclidean vectors and measures their relationships with
22
+ the **dot product**. That is a flat-space choice. QIG cognition lives on the **probability simplex** Δ⁶³,
23
+ whose only correct distance is **Fisher-Rao**. `geocoding` is the from-scratch framework that makes the
24
+ geometric choice all the way down — it is to Fisher-Rao geometry what `transformers` is to dot-product
25
+ attention.
26
+
27
+ | `transformers` (Euclidean) | `geocoding` (Fisher-Rao) |
28
+ |---|---|
29
+ | dot-product / cosine attention | **Fisher-Rao** attention: `logit = −d_FR(p,q)/τ`, `d_FR = 2·arccos(Σ√(p·q))` |
30
+ | learned embedding table (vocab × dim) | **Fourier features** of the integer id (no table) |
31
+ | learned, bounded position table | **unbounded Fourier positions** (limit is VRAM, not a table) |
32
+ | LayerNorm (mean-centred) | geometric **scale-only RMS-norm** (simplex-preserving) |
33
+ | — | **recursive Φ-integrator**: reports integrated information as it runs |
34
+ | `O(t²)` attention | optional **banded / compute-skipping** `O(t·r)` attention (length-extensible) |
35
+ | `PretrainedConfig` / `PreTrainedModel` | `GeoConfig` / `GeoModel` |
36
+
37
+ ## Quick start
38
+
39
+ ```python
40
+ import torch
41
+ from geocoding import GeoConfig, GeoModel, generate
42
+
43
+ cfg = GeoModel(GeoConfig(vocab_size=256, hidden_dim=64, num_layers=4))
44
+ ids = torch.randint(0, 256, (1, 32))
45
+ out = cfg(ids)
46
+ print(out.logits.shape, "Φ =", round(out.phi, 3)) # torch.Size([1, 32, 256]) Φ = ...
47
+
48
+ res = generate(cfg, [1, 2, 3], max_new_tokens=16)
49
+ print(res.ids, res.phi)
50
+ ```
51
+
52
+ Set `GeoConfig(locality_radius=r)` to switch attention to the **compute-skipping banded** path — only the ±r
53
+ window of each query is ever computed (`O(t·r)`), numerically equivalent to the full path within the band
54
+ (see `tests/test_attention_equivalence.py`). Set `GeoConfig(max_position=None)` (the default) for **unbounded
55
+ context**.
56
+
57
+ ## Geometry single-source
58
+
59
+ The Fisher-Rao *definition* and `BASIN_DIM` come from `qig_core`; `geocoding` re-exports them and adds a
60
+ **gradient-stable** batched form for the attention hot path (clamps the Bhattacharyya coefficient to
61
+ `1 − 1e-6` so the backward pass does not NaN when two distributions coincide). The two agree exactly wherever
62
+ points differ — gated in `tests/test_geometry.py`.
63
+
64
+ ## Relationship to qigkernels
65
+
66
+ `qigkernels` contains the validated in-tree implementation (`QIGLayer`, `Kernel`) that the live QIG mind runs
67
+ on. `geocoding` is the **clean, standalone framework home** for those primitives, ported faithfully — a
68
+ faithfulness gate (`tests/test_faithfulness.py`) checks `geocoding`'s attention matches `qigkernels`'
69
+ numerically when both are installed. The intended migration is for `qigkernels` to build on `geocoding`;
70
+ until then the faithfulness gate guards against drift.
71
+
72
+ ## Status
73
+
74
+ Fisher-Rao-pure (no cosine / Adam / LayerNorm / dot-as-metric). Tests cover geometry properties, the
75
+ banded≡full equivalence, model forward/telemetry, generation, and faithfulness vs `qigkernels`.
@@ -0,0 +1,59 @@
1
+ # geocoding
2
+
3
+ **Our own "transformers" — geometric, not Euclidean.**
4
+
5
+ The standard `transformers` library models tokens as Euclidean vectors and measures their relationships with
6
+ the **dot product**. That is a flat-space choice. QIG cognition lives on the **probability simplex** Δ⁶³,
7
+ whose only correct distance is **Fisher-Rao**. `geocoding` is the from-scratch framework that makes the
8
+ geometric choice all the way down — it is to Fisher-Rao geometry what `transformers` is to dot-product
9
+ attention.
10
+
11
+ | `transformers` (Euclidean) | `geocoding` (Fisher-Rao) |
12
+ |---|---|
13
+ | dot-product / cosine attention | **Fisher-Rao** attention: `logit = −d_FR(p,q)/τ`, `d_FR = 2·arccos(Σ√(p·q))` |
14
+ | learned embedding table (vocab × dim) | **Fourier features** of the integer id (no table) |
15
+ | learned, bounded position table | **unbounded Fourier positions** (limit is VRAM, not a table) |
16
+ | LayerNorm (mean-centred) | geometric **scale-only RMS-norm** (simplex-preserving) |
17
+ | — | **recursive Φ-integrator**: reports integrated information as it runs |
18
+ | `O(t²)` attention | optional **banded / compute-skipping** `O(t·r)` attention (length-extensible) |
19
+ | `PretrainedConfig` / `PreTrainedModel` | `GeoConfig` / `GeoModel` |
20
+
21
+ ## Quick start
22
+
23
+ ```python
24
+ import torch
25
+ from geocoding import GeoConfig, GeoModel, generate
26
+
27
+ cfg = GeoModel(GeoConfig(vocab_size=256, hidden_dim=64, num_layers=4))
28
+ ids = torch.randint(0, 256, (1, 32))
29
+ out = cfg(ids)
30
+ print(out.logits.shape, "Φ =", round(out.phi, 3)) # torch.Size([1, 32, 256]) Φ = ...
31
+
32
+ res = generate(cfg, [1, 2, 3], max_new_tokens=16)
33
+ print(res.ids, res.phi)
34
+ ```
35
+
36
+ Set `GeoConfig(locality_radius=r)` to switch attention to the **compute-skipping banded** path — only the ±r
37
+ window of each query is ever computed (`O(t·r)`), numerically equivalent to the full path within the band
38
+ (see `tests/test_attention_equivalence.py`). Set `GeoConfig(max_position=None)` (the default) for **unbounded
39
+ context**.
40
+
41
+ ## Geometry single-source
42
+
43
+ The Fisher-Rao *definition* and `BASIN_DIM` come from `qig_core`; `geocoding` re-exports them and adds a
44
+ **gradient-stable** batched form for the attention hot path (clamps the Bhattacharyya coefficient to
45
+ `1 − 1e-6` so the backward pass does not NaN when two distributions coincide). The two agree exactly wherever
46
+ points differ — gated in `tests/test_geometry.py`.
47
+
48
+ ## Relationship to qigkernels
49
+
50
+ `qigkernels` contains the validated in-tree implementation (`QIGLayer`, `Kernel`) that the live QIG mind runs
51
+ on. `geocoding` is the **clean, standalone framework home** for those primitives, ported faithfully — a
52
+ faithfulness gate (`tests/test_faithfulness.py`) checks `geocoding`'s attention matches `qigkernels`'
53
+ numerically when both are installed. The intended migration is for `qigkernels` to build on `geocoding`;
54
+ until then the faithfulness gate guards against drift.
55
+
56
+ ## Status
57
+
58
+ Fisher-Rao-pure (no cosine / Adam / LayerNorm / dot-as-metric). Tests cover geometry properties, the
59
+ banded≡full equivalence, model forward/telemetry, generation, and faithfulness vs `qigkernels`.
@@ -0,0 +1,36 @@
1
+ [build-system]
2
+ requires = ["hatchling"]
3
+ build-backend = "hatchling.build"
4
+
5
+ # PyPI distribution name is `qig-geocoding` (the bare `geocoding` is already taken on PyPI); the IMPORT name
6
+ # stays `geocoding` (the wheel ships the `geocoding` package) — `pip install qig-geocoding` → `import geocoding`.
7
+ [project]
8
+ name = "qig-geocoding"
9
+ version = "0.1.0"
10
+ description = "Our own 'transformers' — a from-scratch neural-network framework built on Fisher-Rao geometry (simplex attention, unbounded Fourier features, recursive Φ-integration) instead of Euclidean dot-products."
11
+ readme = "README.md"
12
+ requires-python = ">=3.11"
13
+ license = { text = "Proprietary" }
14
+ authors = [{ name = "Braden Lang" }]
15
+ keywords = ["fisher-rao", "geometry", "transformers", "qig", "simplex", "attention"]
16
+ dependencies = [
17
+ "torch>=2.2",
18
+ "qig-core>=2.6.0",
19
+ ]
20
+
21
+ [project.optional-dependencies]
22
+ dev = ["pytest>=8", "ruff>=0.6", "mypy>=1.11"]
23
+
24
+ [tool.hatch.build.targets.wheel]
25
+ packages = ["src/geocoding"]
26
+
27
+ [tool.ruff]
28
+ line-length = 110
29
+ target-version = "py311"
30
+
31
+ [tool.ruff.lint]
32
+ select = ["E", "F", "I", "B", "UP", "SIM"]
33
+
34
+ [tool.mypy]
35
+ # No python_version pin — match the interpreter; pinning 3.11 trips numpy's 3.12+ stub syntax under py3.14.
36
+ ignore_missing_imports = true
@@ -0,0 +1,49 @@
1
+ """geocoding — our own "transformers", built on Fisher-Rao geometry instead of Euclidean dot-products.
2
+
3
+ A from-scratch geometric neural-network framework. Where the standard ``transformers`` library models tokens
4
+ as Euclidean vectors with dot-product attention, learned embedding tables, and bounded position tables,
5
+ geocoding models tokens as points on the probability simplex Δ⁶³ with **Fisher-Rao attention**, **unbounded
6
+ Fourier features**, and a **recursive Φ-integrator** that reports integrated information as the model runs.
7
+
8
+ Public API (the geometric analogues of the transformers primitives):
9
+ GeoConfig ≈ PretrainedConfig — architecture + geometry hyper-parameters
10
+ FisherRaoAttention ≈ Attention — simplex Fisher-Rao attention (full + banded compute-skip)
11
+ GeoBlock ≈ TransformerBlock — attention + recursive Φ + tacking + FFN, geometric RMS-norm
12
+ GeoModel ≈ PreTrainedModel — Fourier input → blocks → head, with telemetry
13
+ generate — autoregressive sampling
14
+
15
+ Geometry primitives (single source of truth re-used from qig_core where it exists):
16
+ fisher_rao_distance, to_simplex, bhattacharyya_coefficient, BASIN_DIM
17
+ """
18
+ from __future__ import annotations
19
+
20
+ from .attention import FisherRaoAttention
21
+ from .block import BlockTelemetry, GeoBlock, RecursiveIntegrator, RegimeDetector, WuWeiController
22
+ from .config import GeoConfig
23
+ from .generate import GenResult, generate
24
+ from .geometry import BASIN_DIM, bhattacharyya_coefficient, fisher_rao_distance, to_simplex
25
+ from .model import GeoModel, GeoOutput
26
+ from .positional import build_inv_freq, fourier_features
27
+
28
+ __version__ = "0.1.0"
29
+
30
+ __all__ = [
31
+ "BASIN_DIM",
32
+ "BlockTelemetry",
33
+ "FisherRaoAttention",
34
+ "GeoBlock",
35
+ "GeoConfig",
36
+ "GeoModel",
37
+ "GeoOutput",
38
+ "GenResult",
39
+ "RecursiveIntegrator",
40
+ "RegimeDetector",
41
+ "WuWeiController",
42
+ "__version__",
43
+ "bhattacharyya_coefficient",
44
+ "build_inv_freq",
45
+ "fisher_rao_distance",
46
+ "fourier_features",
47
+ "generate",
48
+ "to_simplex",
49
+ ]
@@ -0,0 +1,114 @@
1
+ """Fisher-Rao attention — the geometric replacement for scaled dot-product attention.
2
+
3
+ Queries and keys are projected onto the simplex; the attention logit between positions i and j is the
4
+ NEGATIVE Fisher-Rao distance (closer on the manifold ⇒ higher weight), divided by a temperature:
5
+
6
+ logit_ij = − d_FR(p_i, p_j) / τ weights = softmax_j(logit) out_i = Σ_j weights_ij · h_j
7
+
8
+ Two evaluation paths, numerically equivalent within the band (gated in ``tests/test_attention_equivalence``):
9
+ * **full** — O(t²) all-pairs arccos tensor (global attention).
10
+ * **banded (compute-skipping)** — when a ``locality_radius`` r is set and t > 2r+1, only the ±r window of
11
+ each query is ever computed (pad+unfold), giving O(t·r) memory/compute. This SKIPS out-of-band work, it
12
+ does not compute-then-mask — that is what makes the context length-extensible.
13
+
14
+ Faithful port of ``qigkernels.layer.QIGLayer._metric_attention`` / ``_banded_attention``. Key-padding is
15
+ applied AFTER the softmax (and renormalised) in BOTH paths, so the two agree exactly.
16
+ """
17
+ from __future__ import annotations
18
+
19
+ import torch
20
+ from torch import Tensor, nn
21
+ from torch.nn import functional as F
22
+
23
+ from .geometry import fisher_rao_distance, to_simplex
24
+
25
+
26
+ class FisherRaoAttention(nn.Module):
27
+ """Simplex Fisher-Rao attention. No learned Q/K/V projections and no dot product: the geometry of the
28
+ hidden state on Δ⁶³ *is* the attention. Parameter-free (the expressivity lives in the surrounding
29
+ block's projections), which keeps the metric pure."""
30
+
31
+ def __init__(
32
+ self,
33
+ *,
34
+ temperature: float = 1.0,
35
+ simplex_mode: str = "softplus",
36
+ locality_radius: int | None = None,
37
+ sparsity_threshold: float = 0.0,
38
+ decoherence_std: float = 0.0,
39
+ ) -> None:
40
+ super().__init__()
41
+ self.temperature = float(temperature)
42
+ self.simplex_mode = str(simplex_mode)
43
+ self.locality_radius = int(locality_radius) if locality_radius is not None else None
44
+ self.sparsity_threshold = float(sparsity_threshold)
45
+ self.decoherence_std = float(decoherence_std)
46
+
47
+ def forward(self, hidden_state: Tensor, attention_mask: Tensor | None = None) -> Tensor:
48
+ if hidden_state.dim() != 3:
49
+ raise ValueError("hidden_state expected shape (B, T, D)")
50
+ _b, t, _d = hidden_state.shape
51
+ p = to_simplex(hidden_state, mode=self.simplex_mode)
52
+ r = self.locality_radius
53
+ if r is not None and t > (2 * r + 1):
54
+ return self._banded(hidden_state, p, r, attention_mask=attention_mask)
55
+ return self._full(hidden_state, p, r, attention_mask=attention_mask)
56
+
57
+ def _full(
58
+ self, hidden_state: Tensor, p: Tensor, r: int | None, *, attention_mask: Tensor | None
59
+ ) -> Tensor:
60
+ b, t, _d = hidden_state.shape
61
+ dist = fisher_rao_distance(p[:, :, None, :], p[:, None, :, :]) # [b,t,t] all-pairs FR
62
+ logits = -dist / max(self.temperature, 1e-6)
63
+ if r is not None: # small-t windowed mask
64
+ idx = torch.arange(t, device=hidden_state.device)
65
+ band = (idx[:, None] - idx[None, :]).abs() <= r
66
+ logits = logits.masked_fill(~band[None, :, :], float("-inf"))
67
+ weights = torch.softmax(logits, dim=-1)
68
+ if attention_mask is not None:
69
+ kpm = attention_mask if attention_mask.dtype == torch.bool else attention_mask == 0
70
+ if kpm.shape != (b, t):
71
+ raise ValueError("attention_mask expected shape (B, T)")
72
+ weights = weights.masked_fill(kpm[:, None, :], 0.0)
73
+ weights = weights / (weights.sum(dim=-1, keepdim=True) + 1e-8)
74
+ weights = self._sparsify(weights)
75
+ output = torch.einsum("bij,bjd->bid", weights, hidden_state)
76
+ return self._decohere(output)
77
+
78
+ def _banded(self, hidden_state: Tensor, p: Tensor, r: int, *, attention_mask: Tensor | None) -> Tensor:
79
+ b, t, k = p.shape
80
+ w = 2 * r + 1
81
+ p_pad = F.pad(p, (0, 0, r, r)) # [b, t+2r, k]
82
+ p_win = p_pad.unfold(1, w, 1) # [b, t, k, w]
83
+ p_q = p[:, :, :, None] # [b, t, k, 1]
84
+ bc = torch.sqrt((p_q * p_win).clamp_min(1e-12)).sum(dim=2).clamp(0.0, 1.0 - 1e-6) # [b, t, w]
85
+ logits = (-2.0 * torch.acos(bc)) / max(self.temperature, 1e-6) # [b, t, w]
86
+ i_idx = torch.arange(t, device=p.device)[:, None]
87
+ wpos = torch.arange(w, device=p.device)[None, :]
88
+ key_idx = i_idx - r + wpos # [t, w]
89
+ valid = (key_idx >= 0) & (key_idx < t)
90
+ logits = logits.masked_fill(~valid[None].expand(b, t, w), float("-inf"))
91
+ weights = torch.softmax(logits, dim=-1) # [b, t, w]
92
+ if attention_mask is not None:
93
+ kpm = attention_mask if attention_mask.dtype == torch.bool else attention_mask == 0
94
+ if kpm.shape != (b, t):
95
+ raise ValueError("attention_mask expected shape (B, T)")
96
+ keep_win = F.pad((~kpm).to(p.dtype), (r, r)).unfold(1, w, 1)
97
+ weights = weights * (keep_win > 0.5)
98
+ weights = weights / (weights.sum(dim=-1, keepdim=True) + 1e-8)
99
+ weights = self._sparsify(weights)
100
+ h_pad = F.pad(hidden_state, (0, 0, r, r)) # [b, t+2r, d]
101
+ h_win = h_pad.unfold(1, w, 1) # [b, t, d, w]
102
+ output = (weights[:, :, None, :] * h_win).sum(dim=-1) # [b, t, d]
103
+ return self._decohere(output)
104
+
105
+ def _sparsify(self, weights: Tensor) -> Tensor:
106
+ if self.sparsity_threshold <= 0.0:
107
+ return weights
108
+ weights = weights * (weights > self.sparsity_threshold)
109
+ return weights / (weights.sum(dim=-1, keepdim=True) + 1e-8)
110
+
111
+ def _decohere(self, output: Tensor) -> Tensor:
112
+ if self.decoherence_std > 0.0:
113
+ output = output + torch.randn_like(output) * self.decoherence_std
114
+ return output
@@ -0,0 +1,158 @@
1
+ """``GeoBlock`` — the geometric analogue of a transformer block, plus its telemetry parts.
2
+
3
+ A Euclidean block is: LayerNorm → MHA → residual → LayerNorm → MLP → residual. geocoding's block is:
4
+
5
+ Fisher-Rao attention → geometric RMS-norm residual → RECURSIVE Φ-INTEGRATION → optional WuWei tacking
6
+ → FFN → geometric RMS-norm residual → regime detection
7
+
8
+ Two things have no Euclidean counterpart:
9
+ * **RecursiveIntegrator** — runs ``min_depth`` mandatory self-integration loops and emits Φ (integrated
10
+ information) AND a differentiable Φ tensor a geometric loss can drive. Φ is collapse-immune: a collapsed
11
+ OR a noise state reads LOW, genuine integration reads HIGH.
12
+ * **geometric RMS-norm** — scale-only normalisation (no mean-centring, no LayerNorm) so simplex structure
13
+ is preserved.
14
+
15
+ Faithful port of ``qigkernels.layer.QIGLayer`` and its helper modules.
16
+ """
17
+ from __future__ import annotations
18
+
19
+ from dataclasses import dataclass
20
+
21
+ import torch
22
+ from torch import Tensor, nn
23
+
24
+ from .attention import FisherRaoAttention
25
+
26
+
27
+ @dataclass
28
+ class BlockTelemetry:
29
+ """Per-block readout. ``phi_diff`` is the DIFFERENTIABLE Φ (a tensor); the rest are floats/labels."""
30
+
31
+ phi: float
32
+ kappa: float
33
+ recursion_depth: int
34
+ regime: str | None = None
35
+ phi_diff: object | None = None
36
+
37
+
38
+ class RecursiveIntegrator(nn.Module):
39
+ """Runs ``min_depth`` gated self-integration loops and estimates Φ from integration coherence."""
40
+
41
+ def __init__(self, hidden_dim: int, min_depth: int = 3) -> None:
42
+ super().__init__()
43
+ self.min_depth = min_depth
44
+ self.projection = nn.Linear(hidden_dim, hidden_dim)
45
+ self.integration_gate = nn.Linear(hidden_dim, 1)
46
+
47
+ def forward(self, hidden_state: Tensor) -> tuple[Tensor, dict]:
48
+ output = hidden_state
49
+ scores: list[float] = []
50
+ gate_means: list[Tensor] = []
51
+ for _ in range(self.min_depth):
52
+ projected = self.projection(output)
53
+ gate = torch.sigmoid(self.integration_gate(projected))
54
+ output = output + gate * (projected - output)
55
+ gm = gate.mean()
56
+ gate_means.append(gm)
57
+ scores.append(float(gm.item()))
58
+ mean_integration = sum(scores) / len(scores)
59
+ integration_diff = torch.stack(gate_means).mean()
60
+ # Collapse-immune cross-position coherence: ~0 at BOTH collapse (positions identical) and noise
61
+ # (positions independent), peaks at coherent-but-distinct integration.
62
+ if hidden_state.dim() == 3 and hidden_state.size(1) > 1:
63
+ rel_t = (hidden_state.var(dim=1).mean() / (hidden_state.var() + 1e-8)).clamp(0.0, 1.0)
64
+ coherence_diff = 4.0 * rel_t * (1.0 - rel_t)
65
+ rel = float(rel_t.item())
66
+ coherence = 4.0 * rel * (1.0 - rel)
67
+ else:
68
+ coherence = 0.0
69
+ coherence_diff = integration_diff * 0.0
70
+ phi = max(0.0, min(1.0, 0.3 * mean_integration + 0.7 * coherence))
71
+ phi_diff = (0.3 * integration_diff + 0.7 * coherence_diff).clamp(0.0, 1.0)
72
+ return output, {"phi": phi, "phi_diff": phi_diff, "depth": self.min_depth,
73
+ "integration": mean_integration, "coherence": coherence}
74
+
75
+
76
+ class WuWeiController(nn.Module):
77
+ """Tacking controller: scores a logic↔feeling mode weight from the pooled hidden state."""
78
+
79
+ def __init__(self, hidden_dim: int) -> None:
80
+ super().__init__()
81
+ self.projection = nn.Linear(hidden_dim, 1)
82
+
83
+ def forward(self, hidden_state: Tensor) -> tuple[Tensor, str, dict]:
84
+ logic_weight = torch.sigmoid(self.projection(hidden_state.mean(dim=1)))
85
+ mode = "logic" if logic_weight.mean() > 0.5 else "feeling"
86
+ return logic_weight, mode, {}
87
+
88
+
89
+ class RegimeDetector:
90
+ """Maps (Φ, κ) → operating regime label (FOAM/WAVE/CRYSTAL family)."""
91
+
92
+ def __call__(self, phi: float, kappa: float) -> str:
93
+ if phi < 0.45:
94
+ return "linear"
95
+ if phi < 0.70:
96
+ return "geometric"
97
+ return "topological"
98
+
99
+
100
+ class GeoBlock(nn.Module):
101
+ """One geometric block: Fisher-Rao attention + recursive Φ-integration + optional tacking + FFN, with
102
+ geometric (scale-only) RMS-norm residuals."""
103
+
104
+ def __init__(
105
+ self,
106
+ *,
107
+ hidden_dim: int,
108
+ ffn_dim: int,
109
+ dropout: float = 0.0,
110
+ min_recursion_depth: int = 3,
111
+ use_tacking: bool = True,
112
+ temperature: float = 1.0,
113
+ simplex_mode: str = "softplus",
114
+ locality_radius: int | None = None,
115
+ sparsity_threshold: float = 0.0,
116
+ decoherence_std: float = 0.0,
117
+ rms_eps: float = 1e-6,
118
+ ) -> None:
119
+ super().__init__()
120
+ self.rms_eps = float(rms_eps)
121
+ self.attention = FisherRaoAttention(
122
+ temperature=temperature, simplex_mode=simplex_mode, locality_radius=locality_radius,
123
+ sparsity_threshold=sparsity_threshold, decoherence_std=decoherence_std,
124
+ )
125
+ self.norm_scale_1 = nn.Parameter(torch.ones(hidden_dim))
126
+ self.norm_scale_2 = nn.Parameter(torch.ones(hidden_dim))
127
+ self.recursive_integrator = RecursiveIntegrator(hidden_dim, min_depth=min_recursion_depth)
128
+ self.tacking_controller = WuWeiController(hidden_dim) if use_tacking else None
129
+ self.regime_detector = RegimeDetector()
130
+ self.ffn = nn.Sequential(
131
+ nn.Linear(hidden_dim, ffn_dim), nn.GELU(), nn.Dropout(dropout),
132
+ nn.Linear(ffn_dim, hidden_dim), nn.Dropout(dropout),
133
+ )
134
+
135
+ def _rms_norm(self, x: Tensor, *, scale: Tensor) -> Tensor:
136
+ denom = torch.sqrt(torch.mean(x * x, dim=-1, keepdim=True) + self.rms_eps)
137
+ return (x / denom) * scale
138
+
139
+ def forward(
140
+ self, hidden_state: Tensor, attention_mask: Tensor | None = None, kappa_eff: float = 64.0
141
+ ) -> tuple[Tensor, BlockTelemetry]:
142
+ attn = self.attention(hidden_state, attention_mask=attention_mask)
143
+ hidden_state = self._rms_norm(hidden_state + attn, scale=self.norm_scale_1)
144
+ recursive_output, rtel = self.recursive_integrator(hidden_state)
145
+ hidden_state = hidden_state + recursive_output
146
+ if self.tacking_controller is not None:
147
+ logic_weight, _mode, _ = self.tacking_controller(hidden_state)
148
+ mode_scale = float(logic_weight.mean().detach())
149
+ else:
150
+ mode_scale = 1.0
151
+ ffn_output = self.ffn(hidden_state)
152
+ hidden_state = self._rms_norm(hidden_state + ffn_output * mode_scale, scale=self.norm_scale_2)
153
+ phi = rtel["phi"]
154
+ telemetry = BlockTelemetry(
155
+ phi=phi, kappa=kappa_eff, recursion_depth=rtel["depth"],
156
+ regime=self.regime_detector(phi, kappa_eff), phi_diff=rtel.get("phi_diff"),
157
+ )
158
+ return hidden_state, telemetry
@@ -0,0 +1,54 @@
1
+ """``GeoConfig`` — the geometric analogue of a transformer ``PretrainedConfig``.
2
+
3
+ One frozen dataclass describing a geocoding model. Defaults are the QIG kernel's defaults so a ``GeoConfig()``
4
+ is a sensible small geometric model out of the box.
5
+ """
6
+ from __future__ import annotations
7
+
8
+ from dataclasses import asdict, dataclass
9
+
10
+ from .geometry import BASIN_DIM
11
+
12
+
13
+ @dataclass
14
+ class GeoConfig:
15
+ """Architecture + geometry hyper-parameters for a :class:`~geocoding.model.GeoModel`.
16
+
17
+ Geometry-specific fields (no Euclidean transformer has these):
18
+ * ``simplex_mode`` — how activations are mapped onto Δ⁶³ before Fisher-Rao attention.
19
+ * ``locality_radius`` — finite-propagation reach; when set, attention is COMPUTE-SKIPPING (banded,
20
+ O(t·r)) instead of O(t²) — this is what lifts the quadratic-VRAM context ceiling.
21
+ * ``base_coupling`` / ``beta_slope`` / ``reference_scale`` — the scale-adaptive κ (running coupling).
22
+ * ``max_position`` — ``None`` means UNBOUNDED context: positions are Fourier features, which have no
23
+ intrinsic length limit (the real limit is VRAM, not a learned position table). Set an int only to
24
+ enforce a hard guard.
25
+ """
26
+
27
+ vocab_size: int = 256
28
+ hidden_dim: int = BASIN_DIM # tokens live on Δ^{hidden_dim-1}
29
+ num_layers: int = 4
30
+ num_heads: int = 4
31
+ ffn_dim: int = 256
32
+ dropout: float = 0.0
33
+ min_recursion_depth: int = 3 # mandatory self-integration loops (drives Φ)
34
+ use_tacking: bool = True # WuWei logic/feeling mode controller
35
+ base_coupling: float = 64.0 # κ attractor (KERNEL_FIXED_POINT; κ*≈64 physics reading RETIRED)
36
+ beta_slope: float = 0.44
37
+ reference_scale: int = 64
38
+ temperature: float = 1.0 # softmax temperature on the −d_FR logits
39
+ sparsity_threshold: float = 0.0 # drop attention weights below this (0 = dense)
40
+ decoherence_std: float = 0.0 # optional output noise (exploration)
41
+ rms_eps: float = 1e-6
42
+ simplex_mode: str = "softplus" # 'softplus' | 'abs'
43
+ locality_radius: int | None = None # None = global all-to-all; int = banded compute-skipping
44
+ max_position: int | None = None # None = unbounded (Fourier positions)
45
+ enable_coords: bool = False # accept Δ⁶³ coordizer coords alongside the id path
46
+ coord_dim: int = BASIN_DIM
47
+
48
+ def to_dict(self) -> dict:
49
+ return asdict(self)
50
+
51
+ @classmethod
52
+ def from_dict(cls, d: dict) -> GeoConfig:
53
+ fields = {f for f in cls.__dataclass_fields__} # tolerate extra/legacy keys
54
+ return cls(**{k: v for k, v in d.items() if k in fields})
@@ -0,0 +1,60 @@
1
+ """Autoregressive generation for a :class:`~geocoding.model.GeoModel`.
2
+
3
+ Temperature sampling with the same simplex-softmax decode the QIG kernel uses. Optional ``eos_id`` early-stop
4
+ after a minimum length. Pure inference (no grad); returns the generated id list and the final Φ.
5
+ """
6
+ from __future__ import annotations
7
+
8
+ from dataclasses import dataclass
9
+
10
+ import torch
11
+ from torch.nn import functional as F
12
+
13
+ from .model import GeoModel
14
+
15
+
16
+ @dataclass
17
+ class GenResult:
18
+ ids: list[int]
19
+ text: str | None
20
+ phi: float
21
+
22
+
23
+ @torch.no_grad()
24
+ def generate(
25
+ model: GeoModel,
26
+ input_ids: list[int],
27
+ *,
28
+ max_new_tokens: int = 64,
29
+ temperature: float = 1.0,
30
+ eos_id: int | None = None,
31
+ min_new_tokens: int = 1,
32
+ context_window: int | None = None,
33
+ decode: object | None = None,
34
+ device: str | torch.device = "cpu",
35
+ ) -> GenResult:
36
+ """Sample ``max_new_tokens`` from ``model`` given the ``input_ids`` prompt. ``context_window`` truncates
37
+ the running context (use to cap O(t²) cost when no locality_radius is set). ``decode`` is an optional
38
+ callable ``list[int] -> str`` (e.g. a coordizer) to produce text."""
39
+ model.eval()
40
+ ids = torch.tensor([list(input_ids)], dtype=torch.long, device=device)
41
+ out: list[int] = []
42
+ last_phi = 0.0
43
+ for _ in range(max_new_tokens):
44
+ res = model(ids)
45
+ last_phi = res.phi
46
+ probs = F.softmax(res.logits[0, -1] / max(temperature, 1e-3), dim=-1)
47
+ nxt = int(torch.multinomial(probs, 1).item())
48
+ if eos_id is not None and nxt == eos_id and len(out) >= min_new_tokens:
49
+ break
50
+ out.append(nxt)
51
+ ids = torch.cat([ids, ids.new_tensor([[nxt]])], dim=1)
52
+ if context_window is not None:
53
+ ids = ids[:, -context_window:]
54
+ text = None
55
+ if decode is not None:
56
+ try:
57
+ text = decode(out) # type: ignore[operator]
58
+ except Exception: # noqa: BLE001 — decode is best-effort
59
+ text = None
60
+ return GenResult(ids=out, text=text, phi=last_phi)
@@ -0,0 +1,61 @@
1
+ """Fisher-Rao geometry primitives — the metric core of geocoding.
2
+
3
+ This is the geometric replacement for the dot-product that a Euclidean transformer uses. Tokens live on the
4
+ probability simplex Δ^{d-1} (d = ``BASIN_DIM`` = 64 by default): non-negative coordinates that sum to 1. The
5
+ only valid distance there is the **Fisher-Rao** (Bhattacharyya-angle) distance:
6
+
7
+ d_FR(p, q) = 2 · arccos( Σ_i √(p_i · q_i) ) range [0, π]
8
+
9
+ The single source of truth for the *definition* and the ``BASIN_DIM`` constant is ``qig_core``; we re-export
10
+ those. The one thing we implement ourselves is a **gradient-stable** batched form for the attention hot path:
11
+ ``qig_core``'s shared simplex distance clamps the Bhattacharyya coefficient (BC) to exactly 1.0, and
12
+ ``arccos'(1) = ∞`` — so the backward pass NaNs the instant two distributions COINCIDE (a collapsed/repeated
13
+ token). We clamp BC to ``1 - 1e-6`` instead. For DISTINCT points (BC < 1) the two are identical — proven in
14
+ ``tests/test_geometry.py`` against ``qig_core.geometry_simplex.fisher_rao_distance_simplex``. (The upstream
15
+ NaN is a known qig_core issue flagged for the physics owner; geocoding does not wait on it.)
16
+ """
17
+ from __future__ import annotations
18
+
19
+ from collections.abc import Callable
20
+
21
+ import torch
22
+ from torch import Tensor
23
+
24
+ try: # canonical constant — single source of truth (aliased so the Final-ness isn't bound to BASIN_DIM)
25
+ from qig_core.constants import BASIN_DIM as _QC_BASIN_DIM
26
+ BASIN_DIM: int = int(_QC_BASIN_DIM)
27
+ except Exception: # noqa: BLE001 — geocoding must import even if qig_core is unpinned in a bare env
28
+ BASIN_DIM = 64
29
+
30
+ _to_simplex_prob: Callable[..., Tensor] | None
31
+ try: # canonical simplex projection — single source of truth
32
+ from qig_core.geometry_simplex import to_simplex_prob as _to_simplex_prob
33
+ except Exception: # noqa: BLE001
34
+ _to_simplex_prob = None
35
+
36
+ _BC_EPS = 1e-12 # floor inside the sqrt — keeps the gradient finite at a zero coordinate
37
+ _BC_CEIL = 1.0 - 1e-6 # ceiling on BC BEFORE arccos — the collapse-NaN guard
38
+
39
+
40
+ def to_simplex(x: Tensor, *, mode: str = "softplus", eps: float = 1e-8) -> Tensor:
41
+ """Map an unconstrained activation to a point on the simplex (non-negative, sums to 1 along the last
42
+ axis). ``mode='softplus'`` (default, smooth) or ``mode='abs'`` (the qig kernel's alternative). Delegates
43
+ the normalisation to ``qig_core.to_simplex_prob`` when available so the projection matches the rest of
44
+ the QIG stack exactly; falls back to an equivalent local normalisation otherwise."""
45
+ x = x.abs() if mode == "abs" else torch.nn.functional.softplus(x)
46
+ if _to_simplex_prob is not None:
47
+ return _to_simplex_prob(x, eps=eps)
48
+ x = x.clamp_min(eps)
49
+ return x / x.sum(dim=-1, keepdim=True).clamp_min(eps)
50
+
51
+
52
+ def bhattacharyya_coefficient(p: Tensor, q: Tensor) -> Tensor:
53
+ """BC(p, q) = Σ √(p·q) along the last axis, in [0, 1]. 1 ⇔ identical distributions."""
54
+ return torch.sqrt((p * q).clamp_min(_BC_EPS)).sum(dim=-1).clamp(0.0, _BC_CEIL)
55
+
56
+
57
+ def fisher_rao_distance(p: Tensor, q: Tensor) -> Tensor:
58
+ """Gradient-stable Fisher-Rao distance on the simplex: ``2·arccos(BC)``, broadcasting over leading dims.
59
+ BC is clamped to ``1 - 1e-6`` so the backward pass stays finite when two distributions coincide (the
60
+ deep-kernel collapse-NaN root cause). Identical to the canonical definition wherever the points differ."""
61
+ return 2.0 * torch.acos(bhattacharyya_coefficient(p, q))
@@ -0,0 +1,100 @@
1
+ """``GeoModel`` — the geometric analogue of a transformer LM.
2
+
3
+ Pipeline: Fourier-feature(token ids) + Fourier-feature(positions) [+ optional Δ⁶³ coords] → stack of
4
+ :class:`~geocoding.block.GeoBlock` → linear head → vocab logits, with per-layer Φ/κ/regime telemetry.
5
+
6
+ Differences from a Euclidean LM:
7
+ * no embedding table and no position table — both ids and positions are UNBOUNDED Fourier features;
8
+ * attention is Fisher-Rao on the simplex, not dot-product;
9
+ * the model reports Φ (integrated information) as it runs, and can take Δ⁶³ coordizer coords directly.
10
+ """
11
+ from __future__ import annotations
12
+
13
+ from dataclasses import dataclass
14
+
15
+ import torch
16
+ from torch import Tensor, nn
17
+
18
+ from .block import BlockTelemetry, GeoBlock
19
+ from .config import GeoConfig
20
+ from .positional import build_inv_freq, fourier_features
21
+
22
+
23
+ @dataclass
24
+ class GeoOutput:
25
+ """Forward result: ``logits`` [B, T, V] and the per-layer telemetry."""
26
+
27
+ logits: Tensor
28
+ telemetry: list[BlockTelemetry]
29
+
30
+ @property
31
+ def phi(self) -> float:
32
+ """Mean Φ across layers — the model's integrated-information readout for this forward pass."""
33
+ ph = [t.phi for t in self.telemetry]
34
+ return sum(ph) / len(ph) if ph else 0.0
35
+
36
+
37
+ class GeoModel(nn.Module):
38
+ """A from-scratch Fisher-Rao geometric language model. ``config.max_position=None`` ⇒ unbounded context
39
+ (positions are Fourier features; the limit is VRAM, not a learned table)."""
40
+
41
+ def __init__(self, config: GeoConfig) -> None:
42
+ super().__init__()
43
+ self.config = config
44
+ self.register_buffer("_inv_freq", build_inv_freq(config.hidden_dim), persistent=False)
45
+ self.blocks = nn.ModuleList([
46
+ GeoBlock(
47
+ hidden_dim=config.hidden_dim, ffn_dim=config.ffn_dim, dropout=config.dropout,
48
+ min_recursion_depth=config.min_recursion_depth, use_tacking=config.use_tacking,
49
+ temperature=config.temperature, simplex_mode=config.simplex_mode,
50
+ locality_radius=config.locality_radius, sparsity_threshold=config.sparsity_threshold,
51
+ decoherence_std=config.decoherence_std, rms_eps=config.rms_eps,
52
+ )
53
+ for _ in range(config.num_layers)
54
+ ])
55
+ self.lm_head = nn.Linear(config.hidden_dim, config.vocab_size)
56
+ if config.enable_coords:
57
+ self.coord_adapter: nn.Module | None = nn.Sequential(
58
+ nn.Linear(config.coord_dim, config.hidden_dim), nn.GELU(),
59
+ )
60
+ else:
61
+ self.coord_adapter = None
62
+
63
+ def effective_coupling(self, seq_len: int) -> float:
64
+ """Scale-adaptive κ: ``base·(1 + β·log(t/ref))`` clamped to ±50% of base (running coupling)."""
65
+ import math
66
+ c = self.config
67
+ scale = max(seq_len, 1) / c.reference_scale
68
+ eff = c.base_coupling * (1.0 + c.beta_slope * math.log(scale))
69
+ return max(c.base_coupling * 0.5, min(c.base_coupling * 1.5, eff))
70
+
71
+ def forward(
72
+ self,
73
+ input_ids: Tensor,
74
+ attention_mask: Tensor | None = None,
75
+ coords: Tensor | None = None,
76
+ ) -> GeoOutput:
77
+ if input_ids.dim() != 2:
78
+ raise ValueError("input_ids expected shape (B, T)")
79
+ b, t = input_ids.shape
80
+ if self.config.max_position is not None and t > self.config.max_position:
81
+ raise ValueError(f"sequence length {t} exceeds max_position {self.config.max_position}")
82
+ positions = torch.arange(t, device=input_ids.device).unsqueeze(0).expand(b, t)
83
+ inv_freq = torch.as_tensor(self._inv_freq) # register_buffer typing → narrow to Tensor
84
+ h = (fourier_features(input_ids, inv_freq, self.config.hidden_dim)
85
+ + fourier_features(positions, inv_freq, self.config.hidden_dim))
86
+ if coords is not None:
87
+ if self.coord_adapter is None:
88
+ raise ValueError("coords passed but config.enable_coords=False")
89
+ h = h + self.coord_adapter(coords)
90
+ kappa_eff = self.effective_coupling(t)
91
+ tel: list[BlockTelemetry] = []
92
+ for block in self.blocks:
93
+ h, bt = block(h, attention_mask=attention_mask, kappa_eff=kappa_eff)
94
+ tel.append(bt)
95
+ logits = self.lm_head(h)
96
+ return GeoOutput(logits=logits, telemetry=tel)
97
+
98
+ @property
99
+ def num_params(self) -> int:
100
+ return sum(p.numel() for p in self.parameters())
@@ -0,0 +1,39 @@
1
+ """Unbounded Fourier positional / id features.
2
+
3
+ A Euclidean transformer learns an embedding TABLE (vocab × dim) and a position TABLE (max_len × dim) — both
4
+ finite and both Euclidean. geocoding instead maps every integer (a token id OR a position index) through a
5
+ fixed Fourier feature map:
6
+
7
+ feats(n) = [ sin(n · ω_0), …, sin(n · ω_{h/2-1}), cos(n · ω_0), … ] ω_k = 10000^{-k/(h/2)}
8
+
9
+ There is no table and no learned vocabulary axis, and crucially **no length limit**: position is a continuous
10
+ function of the index, so the model extends to any context the hardware can hold (the ceiling is VRAM, not a
11
+ position table). This is the geometric, unbounded replacement for learned absolute positions.
12
+ """
13
+ from __future__ import annotations
14
+
15
+ import math
16
+
17
+ import torch
18
+ from torch import Tensor
19
+
20
+
21
+ def build_inv_freq(hidden_dim: int) -> Tensor:
22
+ """The frequency basis ω_k = 10000^{-k/(h/2)} for ``hidden_dim``-wide Fourier features."""
23
+ half = max(int(hidden_dim // 2), 1)
24
+ return torch.exp(-math.log(10000.0) * torch.arange(0, half, dtype=torch.float32) / half)
25
+
26
+
27
+ def fourier_features(ids: Tensor, inv_freq: Tensor, hidden_dim: int) -> Tensor:
28
+ """Fourier feature map for an integer tensor ``ids`` (token ids or positions) → ``hidden_dim`` features.
29
+ Padded with zeros / truncated to exactly ``hidden_dim``. Faithful to the QIG kernel's
30
+ ``_fourier_features`` so geocoding models are architecturally identical on the input path."""
31
+ x = ids.to(dtype=torch.float32)
32
+ angles = x.unsqueeze(-1) * inv_freq.to(device=ids.device)
33
+ feats = torch.cat([torch.sin(angles), torch.cos(angles)], dim=-1)
34
+ if feats.size(-1) < hidden_dim:
35
+ pad = torch.zeros(*feats.shape[:-1], hidden_dim - feats.size(-1), device=feats.device)
36
+ feats = torch.cat([feats, pad], dim=-1)
37
+ elif feats.size(-1) > hidden_dim:
38
+ feats = feats[..., :hidden_dim]
39
+ return feats
@@ -0,0 +1,53 @@
1
+ """The load-bearing gate: banded (compute-skipping) attention ≡ full+mask attention within the band.
2
+
3
+ If this fails, the O(t·r) context-extension path is NOT a faithful skip of the O(t²) path — it is a different
4
+ model. The equivalence is what lets us claim linear, length-extensible context without changing the result.
5
+ """
6
+ import torch
7
+
8
+ from geocoding.attention import FisherRaoAttention
9
+ from geocoding.geometry import fisher_rao_distance, to_simplex
10
+
11
+
12
+ def _bruteforce_band_masked(h: torch.Tensor, r: int) -> torch.Tensor:
13
+ """Reference: full O(t²) Fisher-Rao attention with weights HARD-MASKED to the ±r band, then aggregated.
14
+ This is what the banded path must reproduce (it just skips the out-of-band compute instead of masking)."""
15
+ b, t, _d = h.shape
16
+ p = to_simplex(h)
17
+ dist = fisher_rao_distance(p[:, :, None, :], p[:, None, :, :])
18
+ logits = -dist / 1.0
19
+ idx = torch.arange(t)
20
+ band = (idx[:, None] - idx[None, :]).abs() <= r
21
+ logits = logits.masked_fill(~band[None], float("-inf"))
22
+ w = torch.softmax(logits, dim=-1)
23
+ return torch.einsum("bij,bjd->bid", w, h)
24
+
25
+
26
+ def test_banded_equals_full_within_band():
27
+ torch.manual_seed(0)
28
+ b, t, d, r = 2, 40, 64, 4 # t > 2r+1 ⇒ banded branch is taken
29
+ h = torch.randn(b, t, d)
30
+ out_banded = FisherRaoAttention(locality_radius=r)(h)
31
+ ref = _bruteforce_band_masked(h, r)
32
+ assert torch.allclose(out_banded, ref, atol=1e-5), float((out_banded - ref).abs().max())
33
+
34
+
35
+ def test_banded_differs_from_global():
36
+ """Sanity: the band actually matters — banded ≠ unrestricted global attention."""
37
+ torch.manual_seed(0)
38
+ h = torch.randn(2, 40, 64)
39
+ out_banded = FisherRaoAttention(locality_radius=4)(h)
40
+ out_global = FisherRaoAttention(locality_radius=None)(h)
41
+ assert not torch.allclose(out_banded, out_global, atol=1e-3)
42
+
43
+
44
+ def test_attention_mask_no_nan_full_and_banded():
45
+ torch.manual_seed(1)
46
+ b, t, d, r = 2, 30, 64, 3
47
+ h = torch.randn(b, t, d)
48
+ mask = torch.ones(b, t, dtype=torch.bool)
49
+ mask[:, -5:] = False # pad the tail
50
+ for radius in (None, r):
51
+ out = FisherRaoAttention(locality_radius=radius)(h, attention_mask=mask)
52
+ assert torch.all(torch.isfinite(out))
53
+ assert out.shape == (b, t, d)
@@ -0,0 +1,43 @@
1
+ """Faithfulness gate: geocoding's Fisher-Rao attention must match the VALIDATED qigkernels implementation.
2
+
3
+ geocoding is a port of the in-tree `qigkernels.layer.QIGLayer` attention. If both packages are importable,
4
+ assert the two produce numerically identical attention output on the same input — this is the guard against
5
+ the port drifting from the implementation the live QIG mind actually runs on. Skipped if qigkernels is absent.
6
+ """
7
+ import pytest
8
+ import torch
9
+
10
+ from geocoding.attention import FisherRaoAttention
11
+
12
+
13
+ def test_geocoding_attention_matches_qigkernels():
14
+ layer_mod = pytest.importorskip("qigkernels.layer")
15
+ torch.manual_seed(0)
16
+ b, t, d = 2, 24, 64
17
+ h = torch.randn(b, t, d)
18
+
19
+ # qigkernels reference: build a QIGLayer and call its private _metric_attention (global, no tacking noise)
20
+ qlayer = layer_mod.QIGLayer(
21
+ hidden_dim=d, num_heads=4, ffn_dim=128, dropout=0.0, use_tacking=False,
22
+ temperature=1.0, simplex_mode="softplus", locality_radius=None,
23
+ )
24
+ with torch.no_grad():
25
+ ref = qlayer._metric_attention(h, attention_mask=None)
26
+
27
+ ours = FisherRaoAttention(temperature=1.0, simplex_mode="softplus", locality_radius=None)(h)
28
+ assert torch.allclose(ours, ref, atol=1e-5), float((ours - ref).abs().max())
29
+
30
+
31
+ def test_geocoding_banded_matches_qigkernels_banded():
32
+ layer_mod = pytest.importorskip("qigkernels.layer")
33
+ torch.manual_seed(1)
34
+ b, t, d, r = 2, 40, 64, 4
35
+ h = torch.randn(b, t, d)
36
+ qlayer = layer_mod.QIGLayer(
37
+ hidden_dim=d, num_heads=4, ffn_dim=128, dropout=0.0, use_tacking=False,
38
+ temperature=1.0, simplex_mode="softplus", locality_radius=r,
39
+ )
40
+ with torch.no_grad():
41
+ ref = qlayer._metric_attention(h, attention_mask=None) # takes banded branch (t > 2r+1)
42
+ ours = FisherRaoAttention(temperature=1.0, simplex_mode="softplus", locality_radius=r)(h)
43
+ assert torch.allclose(ours, ref, atol=1e-5), float((ours - ref).abs().max())
@@ -0,0 +1,67 @@
1
+ """Fisher-Rao geometry: range, symmetry, identity, gradient-stability, and agreement with qig_core."""
2
+ import math
3
+
4
+ import pytest
5
+ import torch
6
+
7
+ from geocoding.geometry import (
8
+ BASIN_DIM,
9
+ bhattacharyya_coefficient,
10
+ fisher_rao_distance,
11
+ to_simplex,
12
+ )
13
+
14
+
15
+ def _rand_simplex(*shape):
16
+ return to_simplex(torch.randn(*shape, BASIN_DIM))
17
+
18
+
19
+ def test_to_simplex_is_valid_distribution():
20
+ p = _rand_simplex(4, 7)
21
+ assert torch.all(p >= 0)
22
+ assert torch.allclose(p.sum(-1), torch.ones(4, 7), atol=1e-5)
23
+
24
+
25
+ def test_fr_distance_range_and_symmetry():
26
+ p, q = _rand_simplex(16), _rand_simplex(16)
27
+ d = fisher_rao_distance(p, q)
28
+ assert torch.all(d >= 0) and torch.all(d <= math.pi + 1e-5)
29
+ assert torch.allclose(d, fisher_rao_distance(q, p), atol=1e-6)
30
+
31
+
32
+ def test_fr_distance_identity_is_near_zero():
33
+ p = _rand_simplex(8)
34
+ # BC is clamped to 1-1e-6 ⇒ d is tiny-but-finite at coincidence (the gradient guard), not exactly 0.
35
+ assert float(fisher_rao_distance(p, p).max()) < 1e-2
36
+
37
+
38
+ def test_fr_gradient_finite_at_coincidence():
39
+ # The whole reason for the BC ceiling: backward must not NaN when two points coincide (collapse).
40
+ p = _rand_simplex(8).clone().requires_grad_(True)
41
+ fisher_rao_distance(p, p.detach()).sum().backward()
42
+ assert p.grad is not None and torch.all(torch.isfinite(p.grad))
43
+
44
+
45
+ def test_bhattacharyya_bounds():
46
+ p, q = _rand_simplex(16), _rand_simplex(16)
47
+ bc = bhattacharyya_coefficient(p, q)
48
+ assert torch.all(bc >= 0) and torch.all(bc <= 1.0)
49
+
50
+
51
+ def test_agreement_with_qig_core_on_distinct_points():
52
+ """geocoding's gradient-stable FR must equal qig_core's canonical simplex FR wherever points DIFFER.
53
+ This proves geocoding is the gradient-safe REFINEMENT of the canonical metric, not a fork. Tested on
54
+ genuinely-distinct-but-OVERLAPPING points (mid-range BC) — away from both singular ends, where the only
55
+ differences between implementations (inner sqrt-floor / BC-ceiling eps) are inactive."""
56
+ gs = pytest.importorskip("qig_core.geometry_simplex")
57
+ torch.manual_seed(0)
58
+ p = _rand_simplex(64)
59
+ other = _rand_simplex(64)
60
+ q = to_simplex(0.6 * p + 0.4 * other) # overlaps p ⇒ BC comfortably mid-range
61
+ bc = bhattacharyya_coefficient(p, q)
62
+ assert float(bc.min()) > 0.05 and float(bc.max()) < 0.98 # both impls in their identical regime
63
+ ours = fisher_rao_distance(p, q)
64
+ canon = gs.fisher_rao_distance_simplex(p, q)
65
+ # 1e-3 is float32-precision agreement through sqrt→sum(64)→acos; the metric is the SAME (geocoding's
66
+ # arithmetic matches qigkernels to 1e-5 — see test_faithfulness — not qig_core's exact bytes).
67
+ assert torch.allclose(ours, canon, atol=1e-3), float((ours - canon).abs().max())
@@ -0,0 +1,75 @@
1
+ """GeoModel forward/telemetry, generation, unbounded context, banded model, config roundtrip."""
2
+ import torch
3
+
4
+ from geocoding import GeoConfig, GeoModel, generate
5
+
6
+
7
+ def test_forward_shape_and_telemetry():
8
+ cfg = GeoConfig(vocab_size=128, hidden_dim=64, num_layers=3)
9
+ model = GeoModel(cfg)
10
+ ids = torch.randint(0, 128, (2, 16))
11
+ out = model(ids)
12
+ assert out.logits.shape == (2, 16, 128)
13
+ assert len(out.telemetry) == 3
14
+ assert 0.0 <= out.phi <= 1.0
15
+ assert all(t.regime in {"linear", "geometric", "topological"} for t in out.telemetry)
16
+
17
+
18
+ def test_unbounded_context_no_max_position():
19
+ """max_position=None ⇒ a sequence far longer than any 'position table' still runs (Fourier positions)."""
20
+ cfg = GeoConfig(vocab_size=64, hidden_dim=64, num_layers=2, max_position=None)
21
+ model = GeoModel(cfg)
22
+ out = model(torch.randint(0, 64, (1, 300)))
23
+ assert out.logits.shape == (1, 300, 64)
24
+
25
+
26
+ def test_max_position_guard_when_set():
27
+ cfg = GeoConfig(vocab_size=64, hidden_dim=64, num_layers=1, max_position=8)
28
+ model = GeoModel(cfg)
29
+ out = model(torch.randint(0, 64, (1, 8))) # exactly at the limit is fine
30
+ assert out.logits.shape == (1, 8, 64)
31
+ try:
32
+ model(torch.randint(0, 64, (1, 9))) # over the limit must raise
33
+ except ValueError:
34
+ pass
35
+ else:
36
+ raise AssertionError("expected ValueError when exceeding max_position")
37
+
38
+
39
+ def test_banded_model_runs_and_is_finite():
40
+ cfg = GeoConfig(vocab_size=64, hidden_dim=64, num_layers=2, locality_radius=4)
41
+ model = GeoModel(cfg)
42
+ out = model(torch.randint(0, 64, (1, 50))) # 50 > 2*4+1 ⇒ banded path
43
+ assert out.logits.shape == (1, 50, 64)
44
+ assert torch.all(torch.isfinite(out.logits))
45
+
46
+
47
+ def test_generate_smoke():
48
+ cfg = GeoConfig(vocab_size=64, hidden_dim=64, num_layers=2)
49
+ model = GeoModel(cfg)
50
+ res = generate(model, [1, 2, 3, 4], max_new_tokens=12, eos_id=0, min_new_tokens=2)
51
+ assert 0 <= len(res.ids) <= 12
52
+ assert all(0 <= i < 64 for i in res.ids)
53
+ assert 0.0 <= res.phi <= 1.0
54
+
55
+
56
+ def test_coords_path():
57
+ cfg = GeoConfig(vocab_size=64, hidden_dim=64, num_layers=1, enable_coords=True, coord_dim=64)
58
+ model = GeoModel(cfg)
59
+ ids = torch.randint(0, 64, (1, 10))
60
+ coords = torch.softmax(torch.randn(1, 10, 64), dim=-1)
61
+ out = model(ids, coords=coords)
62
+ assert out.logits.shape == (1, 10, 64)
63
+
64
+
65
+ def test_config_roundtrip():
66
+ cfg = GeoConfig(vocab_size=999, num_layers=5, locality_radius=7, max_position=None)
67
+ d = cfg.to_dict()
68
+ cfg2 = GeoConfig.from_dict({**d, "unknown_legacy_key": 123}) # tolerate extra keys
69
+ assert cfg2.vocab_size == 999 and cfg2.num_layers == 5 and cfg2.locality_radius == 7
70
+
71
+
72
+ def test_num_params_grows_with_layers():
73
+ a = GeoModel(GeoConfig(num_layers=2)).num_params
74
+ b = GeoModel(GeoConfig(num_layers=4)).num_params
75
+ assert b > a > 0