lgatr.layers.mlp.mlp.GeoMLP

class lgatr.layers.mlp.mlp.GeoMLP(config)

Bases: Module

MLP with geometric product.

This is a core component of L-GATr’s transformer blocks. It is similar to a regular MLP, except that it uses geometric bilinears (the geometric product) in place of the first linear layer.

Assumes input has shape (…, channels, 16), output has shape (…, channels, 16), will create hidden layers with shape (…, increase_hidden_channels*channels, 16).

Parameters:

config (MLPConfig) – Configuration object

forward(multivectors, scalars)

Forward pass.

Parameters:

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

• scalars (None or torch.Tensor) – Optional input scalars with shape (…, s_channels).

Return type:

Tuple[Tensor, Optional[Tensor]]

Returns:

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

• outputs_s (None or torch.Tensor) – Output scalars with shape (…, s_channels).