FullyConnectedTensorProduct

class cuequivariance_torch.FullyConnectedTensorProduct(

irreps_in1: Irreps,

irreps_in2: Irreps,

irreps_out: Irreps,

*,

layout: IrrepsLayout | None = None,

layout_in1: IrrepsLayout | None = None,

layout_in2: 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_in1 (Irreps) – 第一个操作数的输入 irreps。

• irreps_in2 (Irreps) – 第二个操作数的输入 irreps。

• irreps_out (Irreps) – 输出 irreps。

• layout (IrrepsLayout, 可选) – 输入和输出 irreps 的布局。默认为 cue.mul_ir，这是对应于 e3nn 的布局。

• shared_weights (bool, 可选) – 是否跨批次维度共享权重。默认为 True。

• internal_weights (bool, 可选) – 是否为权重创建模块参数。默认为 None。

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

注意

在 e3nn 中，有 irrep_normalization 和 path_normalization 参数。此模块目前仅支持 “component” irrep 归一化和 “element” 路径归一化。

前向传播

forward(

x1: Tensor,

x2: Tensor,

weight: Tensor | None = None,

) → Tensor

执行全连接张量积运算的前向传播。

参数:

• x1 (torch.Tensor) – 第一个操作数的输入张量。它应该具有形状 (batch_size, irreps_in1.dim)。

• x2 (torch.Tensor) – 第二个操作数的输入张量。它应该具有形状 (batch_size, irreps_in2.dim)。

• weight (torch.Tensor, 可选) – 张量积的权重。如果 shared_weights 为 False，则它应该具有形状 (batch_size, weight_numel)；如果 shared_weights 为 True，则它应该具有形状 (weight_numel,)。 如果为 None，则使用内部权重。

返回:

全连接张量积运算产生的输出张量。它将具有形状 (batch_size, irreps_out.dim)。

返回类型:

torch.Tensor

引发:

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