TensorProduct#

class cuequivariance_torch.TensorProduct(
descriptor: SegmentedTensorProduct,
*,
device: device | None = None,
math_dtype: dtype | None = None,
use_fallback: bool | None = None,
)#

PyTorch 模块，用于计算描述符定义的分段张量积的最后一个操作数。

参数:

  • descriptor (SegmentedTensorProduct) – 分段张量积的描述符。

  • math_dtype (torch.dtype, optional) – 系数和计算的数据类型。

  • device (torch.device, optional) – 执行计算的设备。

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

  • Raises – RuntimeError: 如果 use_fallback 为 False 且没有 CUDA 内核可用，则引发此错误。

前向传播

forward(
x0: Tensor,
x1: Tensor | None = None,
x2: Tensor | None = None,
x3: Tensor | None = None,
x4: Tensor | None = None,
x5: Tensor | None = None,
x6: Tensor | None = None,
)#

根据指定的描述符执行张量积。

参数:

  • x0 – 输入张量。输入张量的数量应与描述符中操作数的数量减一相匹配。每个输入张量应具有 (batch, operand_size) 或 (1, operand_size) 的形状，其中 operand_size 对应于张量积描述符中定义的每个操作数的大小。

  • x1[ – 输入张量。输入张量的数量应与描述符中操作数的数量减一相匹配。每个输入张量应具有 (batch, operand_size) 或 (1, operand_size) 的形状，其中 operand_size 对应于张量积描述符中定义的每个操作数的大小。

  • x2 – 输入张量。输入张量的数量应与描述符中操作数的数量减一相匹配。每个输入张量应具有 (batch, operand_size) 或 (1, operand_size) 的形状，其中 operand_size 对应于张量积描述符中定义的每个操作数的大小。

  • x3 – 输入张量。输入张量的数量应与描述符中操作数的数量减一相匹配。每个输入张量应具有 (batch, operand_size) 或 (1, operand_size) 的形状，其中 operand_size 对应于张量积描述符中定义的每个操作数的大小。

  • x4 – 输入张量。输入张量的数量应与描述符中操作数的数量减一相匹配。每个输入张量应具有 (batch, operand_size) 或 (1, operand_size) 的形状，其中 operand_size 对应于张量积描述符中定义的每个操作数的大小。

  • x5 – 输入张量。输入张量的数量应与描述符中操作数的数量减一相匹配。每个输入张量应具有 (batch, operand_size) 或 (1, operand_size) 的形状，其中 operand_size 对应于张量积描述符中定义的每个操作数的大小。

  • x6] – 输入张量。输入张量的数量应与描述符中操作数的数量减一相匹配。每个输入张量应具有 (batch, operand_size) 或 (1, operand_size) 的形状，其中 operand_size 对应于张量积描述符中定义的每个操作数的大小。

返回:

张量积产生的输出张量。它具有 (batch, last_operand_size) 的形状，其中 last_operand_size 是描述符中最后一个操作数的大小。

返回类型:

torch.Tensor

在本页中