lgatr.nets.lgatr_slim.MLP

class lgatr.nets.lgatr_slim.MLP(v_channels, s_channels, nonlinearity='gelu', mlp_ratio=2, num_layers=2, dropout_prob=None)[source]

Bases: Module

Multi-layer perceptron (MLP) for vector and scalar features.

forward(vectors, scalars)[source]
Parameters:

  • vectors (torch.Tensor) – A tensor of shape (…, v_channels, 4) representing Lorentz vectors.

  • scalars (torch.Tensor) – A tensor of shape (…, s_channels) representing scalar features.

Returns:

Tensors of the same shape as input representing the normalized vectors and scalars.

Return type:

torch.Tensor, torch.Tensor