ChannelWiseTensorProduct

class cuequivariance_torch.ChannelWiseTensorProduct(

irreps_in1: Irreps,

irreps_in2: Irreps,

filter_irreps_out: Sequence[Irrep] = None,

*,

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。

• filter_irreps_out (Irrep 序列, 可选) – 输出 irreps 的过滤器。默认为 None。

• 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 张量。