lgatr.layers.slim_layers.SlimLinear
- class lgatr.layers.slim_layers.SlimLinear(in_v_channels, out_v_channels, in_s_channels, out_s_channels, bias=True, initialization='default')[source]
Bases:
ModuleLinear layer for vector and scalar features.
The vector and scalar streams are kept separate; mixing happens elsewhere.
- Parameters:
in_v_channels (
int) – Number of input vector channels.out_v_channels (
int) – Number of output vector channels.in_s_channels (
int) – Number of input scalar channels.out_s_channels (
int) – Number of output scalar channels.bias (
bool) – Whether to include a bias term in the scalar linear layer.initialization (
str) – Initialization scheme for the weights."default"or"small"(smaller weights, used for attention projections to improve stability).
- forward(vectors, scalars)[source]
Apply the linear map.
- Parameters:
vectors (
Tensor) – Lorentz vectors of shape(..., 4, in_v_channels).scalars (
Tensor) – Scalar features of shape(..., in_s_channels).
- Return type:
tuple[Tensor,Tensor]- Returns:
outputs_v – Lorentz vectors of shape
(..., 4, out_v_channels).outputs_s – Scalar features of shape
(..., out_s_channels).