EquivariantTensorProduct#

类 cuequivariance.EquivariantTensorProduct(
d: SegmentedTensorProduct | Sequence[SegmentedTensorProduct],
operands: list[Rep],
)#

等变张量积的描述符。此类是 STP 列表的包装器。虽然 STP 是一个单一的齐次多项式，没有指定每个操作数的角色，但 ETP 确定每个操作数的角色（输入或输出）、每个操作数的表示（irreps）以及每个操作数的布局（多重性优先或不可约表示优先）。

要求

  • 一个 ETP 必须至少包含一个 STP。

  • 每个 STP 必须至少有一个操作数（输出）。

示例

输入0

输入1

输入2

输出

注释

STP0

x

x

x

x

常见情况，操作数数量相同

STP1

x

x

x

某些输入并非所有 STP 都会使用

STP2

x

x

– “ –

STP3

x

– “ –

STP4

x

x

x x x

x

最后一个输入被多次馈送

方法

permute_operands(
permutation: tuple[int, ...],
) → EquivariantTensorProduct#

排列张量积的操作数。

move_operand(
src: int,
dst: int,
) → EquivariantTensorProduct#

将一个操作数移动到新的位置。

move_operand_first(
src: int,
) → EquivariantTensorProduct#

将一个操作数移动到最前面。

move_operand_last(
src: int,
) → EquivariantTensorProduct#

将一个操作数移动到最后面。

squeeze_modes(
modes: str | None = None,
) → EquivariantTensorProduct#

挤压模式。

consolidate_paths() → EquivariantTensorProduct#

合并路径。

canonicalize_subscripts() → EquivariantTensorProduct#

规范化下标。

flatten_modes(
modes: str,
*,
skip_zeros: bool = True,
force: bool = False,
) → EquivariantTensorProduct#

扁平化模式。

all_same_segment_shape() → bool#

是否所有分段都具有相同的形状。

flatten_coefficient_modes() → EquivariantTensorProduct#

扁平化系数模式。

flop_cost(batch_size: int) → int#

计算张量积的浮点运算次数。

memory_cost(
batch_sizes: tuple[int, ...],
itemsize: int | tuple[int, ...],
) → int#

计算张量积的内存访问次数。

backward(
input: int,
) → tuple[EquivariantTensorProduct, tuple[int, ...]]#

等变张量积的反向传播。

参数:

input – 计算反向传播所针对的输入。

返回:

一个元组，包含表示反向传播的 ETP 以及原始 ETP 和反向 ETP 之间操作数的排列。

示例

>>> e = cue.descriptors.fully_connected_tensor_product(
...     cue.Irreps("SO3", "4x0+1x1"),
...     cue.Irreps("SO3", "4x0+2x1"),
...     cue.Irreps("SO3", "4x0+3x1")
... )
>>> e
EquivariantTensorProduct(114x0 x 4x0+1 x 4x0+2x1 -> 4x0+3x1)
>>> e_bwd, i = e.backward(0)
>>> e_bwd
EquivariantTensorProduct(4x0+3x1 x 4x0+1 x 4x0+2x1 -> 114x0)
>>> i
(3, 1, 2, 0)

>>> e = cue.descriptors.spherical_harmonics(cue.SO3(1), [0, 1, 2, 3])
>>> e
EquivariantTensorProduct((1)^(0..3) -> 0+1+2+3)
>>> e_bwd, i = e.backward(0)
>>> e_bwd
EquivariantTensorProduct(0+1+2+3 x (1)^(0..2) -> 1)
>>> i
(1, 0, 0)
classmethod stack(
es: Sequence[EquivariantTensorProduct],
stacked: list[bool],
) → EquivariantTensorProduct#

堆叠多个等变张量积。

symmetrize_operands() → EquivariantTensorProduct#

对 ETP 的操作数进行对称化。

sort_indices_for_identical_operands() → EquivariantTensorProduct#

对相同操作数的索引进行排序。