spd_learn.modules.SPDBatchNormCholesky#

class spd_learn.modules.SPDBatchNormCholesky(num_features: int, momentum: float = 0.1, rebias: bool = True, eps: float = 1e-06, device: device | None = None, dtype: dtype | None = None)[source]#

Bases: Module

Log-Cholesky Riemannian Batch Normalization for SPD Matrices.

Implements Cholesky Batch Normalization (CBN) [Wang et al., 2025].

Parameters:

  • num_features (int) – The size of the SPD matrices (number of features).

  • momentum (float, default=0.1) – Momentum factor for moving average of running mean and variance.

  • rebias (bool, default=True) – If True, applies learnable affine scale and bias in Log-Cholesky space.

  • eps (float, default=1e-6) – Numerical stability epsilon.

  • device (Optional[torch.device], default=None) – Device on which to allocate parameters and buffers.

  • dtype (Optional[torch.dtype], default=None) – Data type of parameters and buffers.

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

Forward pass of Log-Cholesky Riemannian Batch Normalization.

reset_parameters() → None[source]#
reset_running_stats() → None[source]#
running_mean: Tensor#
running_var: Tensor#