equiadapt.nbody.canonicalization_networks package

Submodules

equiadapt.nbody.canonicalization_networks.custom_equivariant_networks module

equiadapt.nbody.canonicalization_networks.custom_group_equivariant_layers module

class equiadapt.nbody.canonicalization_networks.custom_group_equivariant_layers.VNLeakyReLU(in_channels: int, share_nonlinearity: bool = False, negative_slope: float = 0.2)[source]

Bases: Module

Custom module implementing the VN LeakyReLU layer.

forward(x: Tensor) → Tensor[source]

Forward pass of the VNLeakyReLU layer.

Parameters:

x (torch.Tensor) – Input tensor.

Returns:

Output tensor.

Return type:

torch.Tensor

class equiadapt.nbody.canonicalization_networks.custom_group_equivariant_layers.VNSoftplus(in_channels: int, share_nonlinearity: bool = False, negative_slope: float = 0.0)[source]

Bases: Module

Custom module implementing the VN Softplus layer.

forward(x: Tensor) → Tensor[source]

Forward pass of the VNSoftplus layer.

Parameters:

x (torch.Tensor) – Input tensor.

Returns:

Output tensor.

Return type:

torch.Tensor

Module contents