spd_learn.modules.SPDLinear#

class spd_learn.modules.SPDLinear(in_features: int, out_features: int, metric: str = 'LEM', theta: float = 1.0, input_tangent: bool = False, return_tangent: bool = False, device: device | None = None, dtype: dtype | None = None)[source]#

Bases: Module

Intrinsic SPD Transformation Layer.

Implements the intrinsic SPD transformation layer for Riemannian neural networks introduced in RieTrans [Chen et al., 2026].

This layer transforms input SPD matrices in \(\mathcal{S}_{++}^{d_{\text{in}}}\) to output SPD matrices in \(\mathcal{S}_{++}^{d_{\text{out}}}\) (or symmetric tangent matrices in \(\mathcal{S}^{d_{\text{out}}}\)). It leverages an internal SPDRMLR module to project matrices into an unconstrained vectorized space of dimension \(d_{\text{out}}(d_{\text{out}}+1)/2\), reconstructs the symmetric matrix, and optionally applies the matrix exponential to guarantee positive definiteness.

Parameters:

  • in_features (int) – Dimension of the input SPD matrices (\(d_{\text{in}}\)).

  • out_features (int) – Dimension of the output SPD matrices (\(d_{\text{out}}\)).

  • metric ({"LEM", "LCM", "AIM", "PEM"}, default="LEM") – Riemannian metric for the tangent space projection.

  • theta (float, default=1.0) – Deformation parameter for the “PEM” metric.

  • input_tangent (bool, default=False) – If True, expects input to already reside in the tangent space.

  • return_tangent (bool, default=False) – If True, returns symmetric tangent matrices in \(\mathcal{S}^{d_{\text{out}}}\). If False, applies matrix exponential to return strictly SPD matrices.

  • device (torch.device, optional) – Device on which to allocate parameters.

  • dtype (torch.dtype, optional) – Data type of parameters.

References

[1]

Ziheng Chen, Rui Wang, and Nicu Sebe. Building transformation layers for riemannian neural networks. In Advances in Neural Information Processing Systems, volume 39. 2026. URL: https://openreview.net/forum?id=GitZH-Chen/RieTrans.

Examples

>>> import torch
>>> from spd_learn.modules import SPDLinear
>>> layer = SPDLinear(in_features=4, out_features=3, metric="LEM")
>>> X = torch.randn(8, 4, 4)
>>> X = X @ X.mT + 0.1 * torch.eye(4)
>>> Y = layer(X)
>>> Y.shape
torch.Size([8, 3, 3])
extra_repr() → str[source]#

Return the extra representation of the module.

To print customized extra information, you should re-implement this method in your own modules. Both single-line and multi-line strings are acceptable.

forward(X: Tensor) → Tensor[source]#

Forward pass of SPDLinear.

Parameters:

X (torch.Tensor) – Input SPD tensor of shape (…, in_features, in_features) (or tangent tensor if input_tangent=True).

Returns:

Transformed SPD tensor of shape (…, out_features, out_features), or tangent matrix if return_tangent=True.

Return type:

torch.Tensor