spd_learn.functional.gbw_mean#

spd_learn.functional.gbw_mean(values, weights, metric, power=1.0, *, n_iter=10, method='fixed_point')[source]#

Aggregate SPD values using convex attention weights.

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.

  • metric (torch.Tensor) – Common SPD metric; singleton axes must broadcast over keys and queries.

  • power (float or torch.Tensor, default=1.0) – Positive scalar deformation exponent.

  • 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

Notes

Fixed-point updates approximate the barycenter without a convergence guarantee. Float32 accumulation uses float64 internally in the current implementation; root aggregation retains its input dtype.