spd_learn.functional.bw_attention_mean#

spd_learn.functional.bw_attention_mean(values, weights, *, 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.

  • 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.