lgatr.layers.mlp.geometric_bilinears.GeometricBilinear

class lgatr.layers.mlp.geometric_bilinears.GeometricBilinear(in_mv_channels, out_mv_channels, hidden_mv_channels=None, in_s_channels=None, out_s_channels=None)

Bases: Module

Geometric product on multivectors.

Pin-equivariant map between multivector tensors that constructs new geometric features via geometric products.

Parameters:

• in_mv_channels (int) – Input multivector channels

• out_mv_channels (int) – Output multivector channels

• hidden_mv_channels (int or None) – Hidden multivector channels. If None, uses out_mv_channels.

• in_s_channels (int or None) – Input scalar channels. If None, no scalars are expected nor returned.

• out_s_channels (int or None) – Output scalar channels. If None, no scalars are expected nor returned.

forward(multivectors, scalars)

Forward pass.

Parameters:

• multivectors (torch.Tensor) – Input multivectors with shape (…, in_mv_channels, 16)

• scalars (torch.Tensor with shape (..., in_s_channels)) – Input scalars

Return type:

Tuple[Tensor, Tensor]

Returns:

• outputs_mv (torch.Tensor) – Output multivectors with shape (…, out_mv_channels, 16)

• output_s (None or torch.Tensor) – Output scalars with shape (…, out_s_channels).