spd_learn.init.stiefel_#

spd_learn.init.stiefel_(tensor: Tensor, seed: int | None = None) → Tensor[source]#

Initialize tensor on the Stiefel manifold (in-place).

Projects a random matrix to orthonormal columns via polar decomposition. The Stiefel manifold \(\text{St}(n, k)\) is the set of matrices \(W \in \mathbb{R}^{n \times k}\) with orthonormal columns: \(W^\top W = I_k\).

The orthogonal polar factor is computed as:

\[W_{\text{Stiefel}} = W (W^\top W)^{-1/2}\]

Parameters:

tensor (torch.Tensor) – Tensor to initialize, with shape (…, n, k) where n >= k.
seed (int, optional) – Random seed for reproducibility. If None, a warning is logged and seed 0 is used.

Returns:

The initialized tensor (same object as input, modified in-place).

Return type:

torch.Tensor

Notes

Uses dtype-aware eigenvalue clamping from the unified numerical configuration to ensure numerical stability during initialization.

If the eigendecomposition fails (rare), falls back to QR decomposition.

Examples

>>> import torch
>>> from spd_learn import init as spd_init
>>> W = torch.empty(10, 5)
>>> spd_init.stiefel_(W, seed=42)
>>> # Verify orthonormality
>>> torch.allclose(W.T @ W, torch.eye(5), atol=1e-5)
True