feat: depend on pyriemann 0.12 (Array API) for SPD geometry#30
Conversation
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.
There was a problem hiding this comment.
💡 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".
| 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) |
There was a problem hiding this comment.
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 👍 / 👎.
| # 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", |
There was a problem hiding this comment.
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 👍 / 👎.
📚 Documentation Preview📦 Download Documentation Artifact
💡 To enable live previews, add a |
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.Delegated to pyriemann
Thin wrappers, verified value-equivalent to the previous implementations:
airm/log_euclidean/log_cholesky/bures_wasserstein_geodesic→geodesic_*parallel_transport_{airm,lem,log_cholesky}→transport_{riemann,logeuclid,logchol}frechet_derivative_{log,exp}→ddlogm/ddexpmKept in spd_learn (and why)
pyriemann's
_matrix_operatoruses torch's raweighbackward (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:matrix_log/exp/sqrt/inv_sqrt/power),airm/log_euclidean/log_cholesky/bures_wasserstein),Re-exported toolkit
spd_learn.functionalnow re-exports ~79 additional pyriemann geometry functions (distances likedistance_kullback/distance_logdet, means likemean_ale/mean_power/mean_harmonic, kernels,ajd, …), documented inapi.rstunder "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
eighbackward, 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
tests/test_pyriemann_comparison.pymigrated topyriemann.geometry.*(theutils.*paths are deprecated, removed in 0.14) + autograd-ceiling guard test.bures_wasserstein_distancekeeps its subgradient-0 guard at distance 0 (A == B).ruffclean acrossspd_learn/andtests/.Also
AttentionManifold/PhaseDelay(importable, but kept out ofmodels.__all__, which the model test harness iterates as instantiable models).tso a scalartstill broadcasts on batched inputs (pyriemann's_check_alphais strict).