spd_learn.modules.GraphBiMap#

class spd_learn.modules.GraphBiMap(n_nodes: int, in_features: int, out_features: int, *, device=None, dtype=None)[source]#

Bases: Module

Apply graph aggregation followed by node-specific bilinear maps.

See [Ju and Guan, 2024].

Parameters:

  • n_nodes (int) – Number of ordered graph nodes.

  • in_features (int) – Input matrix dimension.

  • out_features (int) – Output matrix dimension.

  • device (torch.device or str, optional) – Device for parameters and buffers.

  • dtype (torch.dtype, optional) – Data type for parameters and buffers.

Notes

Nodes are never mixed across trials. Projections use the host Stiefel parametrization; dimension increases use identity padding.

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

Aggregate and project ordered SPD nodes.

Parameters:

  • input (torch.Tensor) – SPD nodes of shape (…, n_nodes, in_features, in_features).

  • propagation (torch.Tensor) – Row-stochastic weights of shape (…, n_nodes, n_nodes).

Returns:

Projected nodes of shape (…, n_nodes, out_features, out_features).

Return type:

torch.Tensor