0.3.0: full Riemannian Adam + audit fixes (C1-C5, A1-A4, H1-H4)

.github/workflows/test.yml

@@ -9,14 +9,18 @@ on:
 jobs:
   test:
     runs-on: ubuntu-latest
+    strategy:
+      fail-fast: false
+      matrix:
+        python-version: ['3.9', '3.11', '3.12']
 
     steps:
     - uses: actions/checkout@v4
 
-    - name: Set up Python
+    - name: Set up Python ${{ matrix.python-version }}
       uses: actions/setup-python@v5
       with:
-        python-version: '3.12'
+        python-version: ${{ matrix.python-version }}
 
     - name: Install dependencies
       run: |
@@ -27,4 +31,30 @@ jobs:
 
     - name: Run tests
       run: |
-        python -m pytest tests/ -v -s
+        python -m pytest tests/ -v
+
+  package:
+    runs-on: ubuntu-latest
+    steps:
+    - uses: actions/checkout@v4
+    - uses: actions/setup-python@v5
+      with:
+        python-version: '3.12'
+    - name: Build the wheel
+      run: |
+        python -m pip install --upgrade pip build
+        python -m build
+    - name: Check single metadata source
+      run: |
+        test ! -f setup.py
+        python - <<'PY'
+        import glob, zipfile
+        wheels = glob.glob('dist/*.whl')
+        assert len(wheels) == 1, wheels
+        with zipfile.ZipFile(wheels[0]) as z:
+            meta = [n for n in z.namelist() if n.endswith('METADATA')][0]
+            text = z.read(meta).decode()
+        assert 'Version: 0.3.0' in text
+        assert 'Requires-Dist: torch>=2.0.0' in text
+        print('wheel metadata OK')
+        PY

.gitignore

@@ -0,0 +1,8 @@
+__pycache__/
+*.py[cod]
+*.egg-info/
+build/
+dist/
+.pytest_cache/
+.venv/
+*.dat

CHANGELOG.md

@@ -1,5 +1,132 @@
 # Changelog
 
+## 0.3.0 — 2026-11-04
+
+Correctness release for hyperbolic training. The headline change is that
+`k > 0` is now a real Riemannian Adam (exponential map, tangent-space first
+moment, parallel transport) instead of Euclidean Adam with gradient scaling
+and a clip. Several APIs change incompatibly; see **Breaking changes**.
+
+### Breaking changes
+
+- **`resume` defaults to `None`.** An existing state directory is never
+  silently reused or silently reinitialized: `None` raises, `True` continues,
+  `False` reinitializes from `params` with a warning.
+- **Checkpoints are directories.** `save_checkpoint`/`load_checkpoint` are
+  directory-only aliases; the primary API is
+  `export_checkpoint(dir)`/`import_checkpoint(dir)`. 0.2.0 manifests with
+  absolute paths are rejected.
+- **State format v2.** 0.2.0 (`format_version=1`) and 0.1.0 (no meta, no `t`)
+  states are migrated on `resume=True`: weights are preserved, moments and
+  per-row counters are reset because the old math is incompatible with
+  tangent-space moments.
+- **`weight_decay_mode` added.** `"l2"` (default) is part of the Euclidean
+  gradient before the Riemannian conversion (classic Adam L2 / Riemannian
+  gradient of `f + wd/2 ||w||^2`); `"decoupled"` is a tangent-space term and
+  equals AdamW when `k = 0`. The old coordinate shrink is gone.
+- **Source tensors are released by default** (`retain_params=False`):
+  `param_groups[...]["params"]` holds zero-length CPU placeholders after
+  initialization, so a full embedding table is not kept alive in VRAM. Use
+  `read_rows()` for weights, or `retain_params=True` for the old behavior.
+- **Global step is persisted** in `dsa_global_meta.json` and in checkpoints;
+  `resume=True` restores it.
+- **Stricter validation:** negative/out-of-range indices, wrong gradient
+  shapes, non-finite gradients, invalid `k`, and unknown update keys now
+  produce explicit errors or warnings.
+
+### Fixed — audit findings
+
+- **C1 — legacy 0.1.0 could not be reinitialized.** `resume=False` now
+  reinitializes even a partial/legacy directory (missing `t`/meta); weights
+  from `params` are written, moments/counters zeroed. `resume=True` migrates
+  a 0.1.0 directory (weights preserved, `m/v/t` reset, meta v2 written) and a
+  0.2.0 directory (`format_version=1`, same reset rule).
+- **C2 — checkpoints were not portable.** `export_checkpoint(dir)` copies
+  `w/m/v/t` into a directory under relative names and writes
+  `checkpoint.json` with sizes and SHA-256 checksums. `import_checkpoint(dir)`
+  verifies shape, size and hash, stages all files, writes a reinitialization
+  marker, then replaces the live state atomically per file. A moved/copied
+  checkpoint imports successfully; a corrupted one is rejected before the
+  live state changes.
+- **C3 — the optimizer retained the full source tensor.** With
+  `retain_params=False` (default), source tensors are replaced by zero-length
+  placeholders after the on-disk state is written; a weak-reference test
+  confirms the tensor can be collected. `retain_params=True` restores the old
+  PyTorch-style retention.
+- **C4 — weight decay was coordinate shrink.** `weight_decay_mode="l2"` now
+  adds `wd * w` to the Euclidean gradient (Riemannian gradient of the L2
+  penalty), and `"decoupled"` adds `wd * w` to the tangent direction. The
+  Euclidean limit matches Adam L2 and AdamW respectively, and the hyperbolic
+  L2 penalty cancels exactly with `g = -wd * w`.
+- **C5 — two packaging metadata sources.** `setup.py` is removed;
+  `pyproject.toml` is the single source (`torch>=2.0.0`, `numpy>=1.21.0`),
+  matching the README badge. `dsa.__version__` is `0.3.0`.
+- **A1 — a failed `shutdown()` could report success on retry.** The writer is
+  now a `RUNNING → DRAINING → STOPPED` / `FAILED` state machine. A timeout
+  leaves the worker alive and the queue unfinished; the next call retries and
+  only a clean drain/join/flush returns success. Persistent write errors are
+  re-raised, never hidden.
+- **A2 — `global_step` was lost on load.** It is stored in
+  `dsa_global_meta.json`, exported in checkpoints, restored on `resume=True`
+  and `import_checkpoint()`.
+- **A3 — negative indices silently updated the last row.** Indices are
+  validated: 1-D integer dtype, `0 <= i < N`; negative, out-of-range, float
+  and 2-D indices raise before any state is touched.
+- **A4 — a NaN row left a corrupted counter/moment.** A non-finite (or
+  negative-variance) row is fully reset (`w=m=v=0, t=0`), so the current step
+  is a genuine first Adam step, not a continuation of a broken counter.
+- **A9 — gradient shape/finiteness was unchecked.** Wrong shapes, non-float
+  dtypes and NaN/inf gradients raise `ValueError` before moments are updated.
+- **H1 — initialization and resume could leave rows outside the ball.**
+  Fresh hyperbolic weights are globally scaled so the largest row norm is at
+  most `init_max_norm_fraction` (default 0.5) of the radius; rows outside are
+  still projected. On resume, finite out-of-ball rows are projected and
+  counted; `read_rows()` always projects. `export_checkpoint()` immediately
+  after initialization is valid.
+- **H2 — the step was not Riemannian Adam.** The step now uses
+  `grad_R = lambda^-2 g_E`, tangent-space moments, the exponential map
+  `exp_w(-lr * m_hat/(sqrt(v_hat)+eps))` and parallel transport of `m` (and
+  of `sqrt(v)`, squared back to keep the diagonal second moment
+  non-negative). Active rows are computed in `float64` (storage stays
+  `float32`) to avoid catastrophic cancellation in the gyration formula near
+  the boundary. Tests compare against an independent float64 reference and
+  check the defining geometric properties (exp map distance, transport
+  identity/isometry/round-trip, log-exp inverse).
+- **H3 — changing `k` left untouched rows outside the new ball.**
+  `set_curvature(k)` and `step(current_k=...)` reproject the whole table;
+  `reset_moments=True` optionally restarts the adaptive state.
+- **H4 — `k` was not validated.** `k = NaN` no longer falls through to the
+  Euclidean path, `k = inf` no longer collapses the ball, and `k < 0` raises
+  at construction instead of a later `math domain error`.
+- **F1–F8, R1–R3** from 0.2.0 remain covered: lossless queue fallback,
+  in-flight read consistency, strict state lifecycle, per-row bias
+  correction, sorted indices, honest shutdown, projection on read, NaN-row
+  handling, reinitialization marker, and unknown-key diagnostics.
+
+### Tests
+
+The suite now has explicit acceptance tests for C1–C5, A1–A4, A9, H1–H4,
+F1–F8 and R1–R3, plus geometric property tests and an independent float64
+reference. Run with a loaded `.env` if your environment requires it:
+
+```bash
+CUDA_VISIBLE_DEVICES="" python -m pytest -q
+```
+
+### Known limitations
+
+- The Poincaré chart is numerically singular at the boundary. Points are kept
+  at least `BALL_EPS = 1e-5` inside, but very long training with a large `lr`
+  can reach that margin; the optimizer stays finite and projects, while the
+  chart itself is ill-conditioned there.
+- `v` is a per-coordinate diagonal moment. Transporting it directly as a
+  tangent vector can produce negative components and NaNs, so its
+  non-negative square root is transported and squared back. This is the
+  positivity-preserving version of the paper's operation, not a
+  component-wise `sqrt` of a rotated vector.
+- CUDA-specific memory behavior is tested only through object lifetime on CPU
+  in CI; no GPU is required for the test suite.
+
 ## 0.2.0 — 2026-11-04
 
 Correctness release. The public API is preserved: the constructor signature