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,
)#

一个表示等变线性层的类。

参数:
  • irreps_in (Irreps) – 输入不可约表示。

  • irreps_out (Irreps) – 输出不可约表示。

  • layout (IrrepsLayout, 可选) – 不可约表示的布局,默认为 cue.mul_ir。这是 e3nn 库中使用的布局。

  • shared_weights (bool, 可选) – 是否使用共享权重,默认为 True。

  • internal_weights (bool, 可选) – 是否使用内部权重,如果 shared_weights 为 True,则默认为 True,否则为 False。

  • use_fallback (bool, 可选) – 如果为 None (默认值),则在可用时使用 CUDA 内核。如果为 False,将使用 CUDA 内核,如果不可用则会引发异常。如果为 True,则无论 CUDA 内核是否可用,都使用 PyTorch 后备方法。

前向传播

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

线性层的前向传播。

参数:
  • x (torch.Tensor) – 输入张量。

  • weight (torch.Tensor, 可选) – 权重张量。如果为 None,则使用内部权重张量。

返回值:

应用线性变换后的输出张量。

返回类型:

torch.Tensor

抛出:

ValueError – 如果使用内部权重且 weight 不为 None,或者如果使用共享权重且 weight 不是 1D 张量,或者如果未使用共享权重且 weight 不是 2D 张量。