spd_learn.modules.GBWGeometry#

class spd_learn.modules.GBWGeometry(n_features, power=1.5, learn_power=False, learn_metric=True)[source]#

Bases: Module

Parametrize a shared SPD metric and positive deformation power.

See [Jin et al., 2026].

Parameters:

  • n_features (int) – SPD matrix dimension. Metric starts at identity and is learned through the existing symmetric matrix exponential parametrization.

  • power (float, default=1.5) – Positive deformation exponent; zero/negative powers are excluded here.

  • learn_power (bool, default=False) – Learn power via softplus. Paper experiments use fixed 1 or 1.5.

  • learn_metric (bool, default=True) – Whether the common metric is learned. False fixes it at identity.

Notes

Metric and power are registered, serializable state, follow .to(), and need no manifold optimizer. Ordinary optimizers operate on unconstrained originals; this differs from Geoopt retractions. No running statistics.

distance(x, y, *, squared=False)[source]#

Compute pairwise/broadcast distances using the current geometry.

Parameters:

  • x (torch.Tensor) – Real SPD matrices of shape (…, n, n).

  • y (torch.Tensor) – Real SPD matrices broadcastable with x.

  • squared (bool, default=False) – Return squared distances instead of distances.

Returns:

Distances over the broadcast leading axes, preserving input dtype/device.

Return type:

torch.Tensor

forward(x)[source]#

Map SPD matrices into generalized BW coordinates.

Parameters:

x (torch.Tensor) – SPD matrices of shape (…, n_features, n_features).

Returns:

Charted matrices with the input shape.

Return type:

torch.Tensor

inverse(x)[source]#

Map charted SPD matrices back to the original coordinates.

Parameters:

x (torch.Tensor) – Charted matrices of shape (…, n_features, n_features).

Returns:

SPD matrices with the input shape.

Return type:

torch.Tensor

mean(values, weights, *, n_iter=10, method='fixed_point')[source]#

Weighted barycenters; see functional gbw_mean for axes/accuracy.

Parameters:

  • values (torch.Tensor) – SPD values of shape (…, keys, n, n).

  • weights (torch.Tensor) – Nonnegative weights of shape (…, queries, keys), summing to one over keys.

  • n_iter (int, default=10) – Positive number of fixed-point updates.

  • method ({'fixed_point', 'root'}, default='fixed_point') – Fixed-point BW iterations or the weighted-root approximation.

Returns:

SPD means of shape (…, queries, n, n).

Return type:

torch.Tensor

metric: Tensor#
power: Tensor#