Linear#

class cuequivariance_torch.Linear(
irreps_in: Irreps,
irreps_out: Irreps,
*,
layout: IrrepsLayout | None = None,
layout_in: IrrepsLayout | None = None,
layout_out: IrrepsLayout | None = None,
shared_weights: bool = True,
internal_weights: bool = None,
device: device | None = None,
dtype: dtype | None = None,
math_dtype: dtype | None = None,
use_fallback: bool | None = None,
)#

A class that represents an equivariant linear layer.

Parameters:

  • irreps_in (Irreps) – The input irreducible representations.

  • irreps_out (Irreps) – The output irreducible representations.

  • layout (IrrepsLayout, optional) – The layout of the irreducible representations, by default cue.mul_ir. This is the layout used in the e3nn library.

  • shared_weights (bool, optional) – Whether to use shared weights, by default True.

  • internal_weights (bool, optional) – Whether to use internal weights, by default True if shared_weights is True, otherwise False.

  • use_fallback (bool, optional) – If None (default), a CUDA kernel will be used if available. If False, a CUDA kernel will be used, and an exception is raised if it’s not available. If True, a PyTorch fallback method is used regardless of CUDA kernel availability.

Forward Pass

forward(
x: Tensor,
weight: Tensor | None = None,
) → Tensor#

Forward pass of the linear layer.

Parameters:

  • x (torch.Tensor) – The input tensor.

  • weight (torch.Tensor, optional) – The weight tensor. If None, the internal weight tensor is used.

Returns:

The output tensor after applying the linear transformation.

Return type:

torch.Tensor

Raises:

ValueError – If internal weights are used and weight is not None, or if shared weights are used and weight is not a 1D tensor, or if shared weights are not used and weight is not a 2D tensor.

On this page