Fix Stabilize Stiefel init and match PyTorch’s generator API

[lucas-heck]

Problem

Stiefel initialization can produce weights that fail orthogonality checks in float32, particularly for square matrices.

It also used a seed argument, unlike the other initialization methods and PyTorch's convention of accepting an optional generator.

Fix

• Compute the polar decomposition and QR fallback in float64, then cast the result back to the input dtype. Use CPU for the decomposition on MPS, because it does not support float64.
• Use PyTorch's default RNG unless an explicit generator is provided.
• Replace stiefel_'s seed argument with generator and remove BiMap's seed argument. Reproducible BiMap initialization now follows torch.manual_seed.

Verification

• Tightened initialization orthogonality checks to 1e-5, added the previous failing CPU float32 case, and enabled the square Stiefel BiMap case.
• Compared training against the main branch on CPU/CUDA, with matched random draws resulted in identical float64 results and float32 differences below 1e-5.

[codecov-commenter]

⚠️ Please install the to ensure uploads and comments are reliably processed by Codecov.

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!

[bruAristimunha]

Thanks @lucas-heck! This is a good cleanup. Dropping the hidden "seed 0 when no seed is given" default and following torch.manual_seed / generator like torch.nn.init is the right call. The batched torch.linalg.qr(..., mode="reduced") fallback also removes the per-matrix loop.

Checked before merging:

• Full suite with all four open PRs merged together: green.
• Apple MPS (the float64-on-CPU path): float32 16×16 batched Stiefel init gives max|WᵀW − I| ≈ 1.8e-7, and the result comes back on mps in float32.
• No remaining seed= callers of BiMap / stiefel_ in spd_learn/, the examples, or the docs.

One note for the release, not blocking: removing seed from BiMap / stiefel_ breaks existing calls, which will now raise a TypeError. For a 0.x release, the entry under "Changed" is enough. I'll mention it in the release notes.

I'll merge main into the branch to resolve the CHANGELOG.md conflict (both entries kept), then merge. Thanks!