lgatr.layers.mlp.geometric_bilinears.GeometricBilinear

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

Bases: Module

Pin-equivariant bilinear map that constructs new geometric features via geometric products.

The geometric-product output is passed through an EquiLayerNorm before it is returned.

Parameters:

• in_mv_channels (int) – Input multivector channels.

• out_mv_channels (int) – Output multivector channels.

• primitives (PrimitivesConfig) – LGATr primitives configuration.

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

• in_s_channels (int) – Input scalar channels. Use 0 for no scalar inputs.

• out_s_channels (int) – Output scalar channels. Use 0 for no scalar outputs.

forward(multivectors, scalars=None)

Forward pass.

Parameters:

• multivectors (Tensor) – Input multivectors of shape (..., in_mv_channels, 16).

• scalars (Tensor | None) – Optional input scalars of shape (..., in_s_channels). If None, the scalar stream is bypassed.

Return type:

tuple[Tensor, Tensor | None]

Returns:

• outputs_mv – Output multivectors of shape (..., out_mv_channels, 16).

• outputs_s – Output scalars of shape (..., out_s_channels), or None if out_s_channels == 0.