lgatr.nets.lgatr_slim.RMSNorm

class lgatr.nets.lgatr_slim.RMSNorm(epsilon=0.01)[source]

Bases: Module

Normalize jointly over vector and scalar features.

For vectors, we use the absolute value of the squared norm because otherwise negative norms are possible.

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