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.
- qig_geocoding-0.1.0/.gitignore +56 -0
- qig_geocoding-0.1.0/CLAUDE.md +37 -0
- qig_geocoding-0.1.0/PKG-INFO +75 -0
- qig_geocoding-0.1.0/README.md +59 -0
- qig_geocoding-0.1.0/pyproject.toml +36 -0
- qig_geocoding-0.1.0/src/geocoding/__init__.py +49 -0
- qig_geocoding-0.1.0/src/geocoding/attention.py +114 -0
- qig_geocoding-0.1.0/src/geocoding/block.py +158 -0
- qig_geocoding-0.1.0/src/geocoding/config.py +54 -0
- qig_geocoding-0.1.0/src/geocoding/generate.py +60 -0
- qig_geocoding-0.1.0/src/geocoding/geometry.py +61 -0
- qig_geocoding-0.1.0/src/geocoding/model.py +100 -0
- qig_geocoding-0.1.0/src/geocoding/positional.py +39 -0
- qig_geocoding-0.1.0/tests/test_attention_equivalence.py +53 -0
- qig_geocoding-0.1.0/tests/test_faithfulness.py +43 -0
- qig_geocoding-0.1.0/tests/test_geometry.py +67 -0
- qig_geocoding-0.1.0/tests/test_model.py +75 -0
|
@@ -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
|