spd_learn.modules.SPDLogTokenizer#

class spd_learn.modules.SPDLogTokenizer(matrix_size=9, *, preserve_norm=False, eigenvalue_floor=None, matrix_scale=1.0, device=None, dtype=None)[source]#

Bases: Module

Convert prepared SPD matrices to symmetric log tokens.

See [Seraphim et al., 2024].

Parameters:

  • matrix_size (int, default=9) – Size of each prepared (already augmented/whitened, if desired) matrix.

  • preserve_norm (bool, default=False) – Raw canonical upper triangle by default, matching the author code. True uses sqrt(2) off-diagonal scaling and the matching inverse.

  • eigenvalue_floor (float or None) – Optional explicit PSD rectification. Core log numerical floor also applies.

  • matrix_scale (float, default=1.0) – Positive multiplier before the log and spectral floor.

  • device (torch.device or str, optional) – Accepted for API consistency; this stateless operation follows the input device.

  • dtype (torch.dtype, optional) – Accepted for API consistency; this stateless operation follows the input dtype.

Notes

The optional eigenvalue floor is explicit regularization. inverse uses the same basis and scale but cannot undo flooring. Dtype and device follow the input.

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

Tokenize prepared SPD matrices.

Parameters:

x (torch.Tensor) – Matrices of shape (…, matrix_size, matrix_size).

Returns:

Log tokens of shape (…, matrix_size * (matrix_size + 1) // 2).

Return type:

torch.Tensor

inverse(tokens: Tensor) → Tensor[source]#

Reconstruct matrices from log tokens.

Parameters:

tokens (torch.Tensor) – Coordinates of shape (…, matrix_size * (matrix_size + 1) // 2).

Returns:

Matrices of shape (…, matrix_size, matrix_size).

Return type:

torch.Tensor