Skip to content

feat: depend on pyriemann 0.12 (Array API) for SPD geometry#30

Merged
bruAristimunha merged 2 commits into
mainfrom
feat/pyriemann-array-api-delegation
Jun 26, 2026
Merged

feat: depend on pyriemann 0.12 (Array API) for SPD geometry#30
bruAristimunha merged 2 commits into
mainfrom
feat/pyriemann-array-api-delegation

Conversation

@bruAristimunha

Copy link
Copy Markdown
Contributor

Summary

Makes pyriemann>=0.12 (its new Python Array API backend, which runs natively on PyTorch tensors with autograd) a core dependency, and reorganizes SPD geometry around it: delegate what need not be numerically special, keep what does, and re-export pyriemann's broader toolkit so the full SPD/Riemannian API is reachable from one namespace.

Dependency pin. The Array API backend lives in pyriemann 0.12, not yet on PyPI, so pyproject.toml hard-pins pyriemann @ git+https://github.com/pyRiemann/pyRiemann.git. Because this is a direct-URL dependency, spd_learn cannot be uploaded to PyPI until pyriemann 0.12 releases — then the pin becomes pyriemann>=0.12. (Noted inline in pyproject.toml and installation.rst.)

Delegated to pyriemann

Thin wrappers, verified value-equivalent to the previous implementations:

  • Geodesicsairm/log_euclidean/log_cholesky/bures_wasserstein_geodesicgeodesic_*
  • Parallel transportparallel_transport_{airm,lem,log_cholesky}transport_{riemann,logeuclid,logchol}
  • Fréchet derivativesfrechet_derivative_{log,exp}ddlogm/ddexpm

Kept in spd_learn (and why)

pyriemann's _matrix_operator uses torch's raw eigh backward (no Loewner stabilization) and does not clamp eigenvalues. Delegating these would return NaN on ill-conditioned float32 inputs and produce unstable/NaN gradients near clustered eigenvalues — breaking SPDNet/TensorCSPNet/LieBN training and complex-dtype batch norm. So these keep their numerically-stable, custom-autograd, eigenvalue-clamped implementations:

  • the 5 core matrix-function primitives (matrix_log/exp/sqrt/inv_sqrt/power),
  • the 4 distances (airm/log_euclidean/log_cholesky/bures_wasserstein),
  • the weighted/iterative means, the AIRM/LEM tangent exp/log maps, the Cholesky maps, and the ladder transports.

Re-exported toolkit

spd_learn.functional now re-exports ~79 additional pyriemann geometry functions (distances like distance_kullback/distance_logdet, means like mean_ale/mean_power/mean_harmonic, kernels, ajd, …), documented in api.rst under "pyRiemann Geometry (re-exported)". The 7 low-level primitives (logm/expm/sqrtm/invsqrtm/powm/ddlogm/ddexpm) are intentionally not re-exported so spd_learn's stable versions win.

Known consequence (accepted, documented)

Gradients of the delegated ops inherit pyriemann's eigh backward, unstable for near-degenerate eigenvalues. Captured in a guard test (test_delegated_gradients_finite_well_conditioned); the stability-sensitive paths (primitives, distances, means) are deliberately kept out of delegation.

Test plan

  • 709 passed, 146 skipped, 0 failed (full suite).
  • tests/test_pyriemann_comparison.py migrated to pyriemann.geometry.* (the utils.* paths are deprecated, removed in 0.14) + autograd-ceiling guard test.
  • Verified the 4 distances stay finite (forward + grad) on ill-conditioned float32, and bures_wasserstein_distance keeps its subgradient-0 guard at distance 0 (A == B).
  • ruff clean across spd_learn/ and tests/.

Also

  • Documented AttentionManifold / PhaseDelay (importable, but kept out of models.__all__, which the model test harness iterates as instantiable models).
  • Geodesic wrappers coerce a 0-d tensor t so a scalar t still broadcasts on batched inputs (pyriemann's _check_alpha is strict).

Make pyriemann>=0.12 a core dependency (git-pinned until 0.12 releases on
PyPI) and delegate the geometry it need not keep numerically special to its
Array API backend, which runs natively on torch tensors:

- Delegated to pyriemann: the 4 geodesics, AIRM/LEM/log-Cholesky parallel
  transport, and Frechet derivatives (ddlogm/ddexpm).
- Re-exported pyriemann's broader geometry toolkit (~79 functions:
  distances/means/geodesics/tangent maps/kernels/ajd) from
  spd_learn.functional so the full SPD/Riemannian API is one import surface.

Kept in spd_learn (numerically-stable, custom-autograd, eigenvalue-clamped),
because pyriemann's unclamped log returns NaN on ill-conditioned float32 and
its raw eigh backward is unstable near clustered eigenvalues (both break
SPDNet/TensorCSPNet/LieBN training and complex-dtype batch norm):

- the 5 core matrix-function primitives (matrix_log/exp/sqrt/inv_sqrt/power),
- the 4 distances (airm/log_euclidean/log_cholesky/bures_wasserstein),
- the weighted/iterative means, tangent exp/log maps, and ladder transports.

Also: modernize the pyriemann comparison test (pyriemann.geometry.* imports,
autograd-ceiling guard test), document AttentionManifold/PhaseDelay, and update
the installation/FAQ/API docs.

@chatgpt-codex-connector chatgpt-codex-connector Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 26e3c05d7a

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

  • Open a pull request for review
  • Mark a draft as ready
  • Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

Comment on lines +187 to +190
if torch.is_tensor(t) and t.ndim == 0 and A.ndim > 2:
t = t.expand(A.shape[:-2])
# Delegated to pyriemann (Array API, runs on torch tensors with autograd).
return geodesic_riemann(A, B, alpha=t)

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Badge Coerce tensor interpolation weights to input device

When callers pass t as a tensor on a different device/dtype than A (for example torch.tensor(0.5) with CUDA SPD batches), this forwards or expands that tensor unchanged. pyriemann expands tensor alphas before doing arithmetic, so the delegated geodesic then combines a CPU alpha with CUDA matrices and raises a device mismatch; the old implementation coerced t to A. Please move/cast tensor t before passing it to the delegated geodesic, and apply the same fix to the other wrappers using this pattern.

Useful? React with 👍 / 👎.

Comment thread pyproject.toml
# Hard-pinned to git until 0.12 is released on PyPI. NOTE: a direct-URL
# dependency blocks uploading spd_learn to PyPI; swap to "pyriemann>=0.12"
# once 0.12 ships.
"pyriemann @ git+https://github.com/pyRiemann/pyRiemann.git",

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Badge Pin pyriemann to an immutable ref

This comment says the dependency is hard-pinned, but the direct URL omits a tag or commit, so every install resolves whatever pyRiemann's default branch points to at that moment. Since core geometry is now imported from pyriemann at module import time, upstream master changes can silently change numerical results or break installs, making experiments unreproducible. Please pin a commit/tag until this can be replaced with pyriemann>=0.12.

Useful? React with 👍 / 👎.

@github-actions

Copy link
Copy Markdown

📚 Documentation Preview

📦 Download Documentation Artifact

Download the documentation-html artifact from the workflow run to view the docs locally.

💡 To enable live previews, add a SURGE_TOKEN secret to this repository. See surge.sh for setup instructions.

@bruAristimunha
bruAristimunha merged commit affbca4 into main Jun 26, 2026
11 checks passed
@bruAristimunha
bruAristimunha deleted the feat/pyriemann-array-api-delegation branch June 26, 2026 10:11
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant