Interaction

We currently recommend using CgtpInteraction. It supports operator fusion via OpenEquivariance or CuEquivariance, which can significantly reduce memory usage and improve computational efficiency.

Although SO2Interaction is theoretically more advantageous at large angular momentum, it currently lacks support for operator fusion libraries, and is therefore not recommended.

class tace.models._e3nn.inter.CgtpInteraction(layer: int, num_layers: int, num_elements: int, avg_num_neighbors: float, mmax: int, Lmax: int, lmax: int, correlation: list[int], num_channel: int, edge_feats_channel: int, target_irreps: list[str], num_radial_basis: int, radial_mlp: list[int], radial_bias: bool, irreps_in: Irreps, l1l2: str | None = None, scatter_norm: str = 'avg_num_neighbors', bias: bool = True, nonlinear: str | None = None, edge_nonlinear: str | None = None, edge_info_type: str = 'mlp', resnet_type: str = 'BB', resnet_linear_type: str = 'aware', use_first_resnet: bool = False, pre_norm_type: str | None = None, use_first_pre_norm: bool = False, so2_angular_basis: SO3Rotation | None = None, use_so2_edge_ace: bool = False, stochastic_depth: float = 0.0, use_first_dropout: bool = False, parity: bool = False, num_head: int | None = None, use_graph_softmax: bool = False, node_wise_hidden: int | None = None, edge_wise_hidden: int | None = None)

Bases: Interaction

An interaction module based on Clebsch-Gordan tensor products (CGTP).

This module performs edge-level convolution using Clebsch-Gordan tensor products. It supports operator fusion via OpenEquivariance or CuEquivariance, which can significantly reduce memory consumption and improve efficiency.

class tace.models._e3nn.inter.SO2Interaction(layer: int, num_layers: int, num_elements: int, avg_num_neighbors: float, mmax: int, Lmax: int, lmax: int, correlation: list[int], num_channel: int, edge_feats_channel: int, target_irreps: list[str], num_radial_basis: int, radial_mlp: list[int], radial_bias: bool, irreps_in: Irreps, l1l2: str | None = None, scatter_norm: str = 'avg_num_neighbors', bias: bool = True, nonlinear: str | None = None, edge_nonlinear: str | None = None, edge_info_type: str = 'mlp', resnet_type: str = 'BB', resnet_linear_type: str = 'aware', use_first_resnet: bool = False, pre_norm_type: str | None = None, use_first_pre_norm: bool = False, so2_angular_basis: SO3Rotation | None = None, use_so2_edge_ace: bool = False, stochastic_depth: float = 0.0, use_first_dropout: bool = False, parity: bool = False, num_head: int | None = None, use_graph_softmax: bool = False, node_wise_hidden: int | None = None, edge_wise_hidden: int | None = None)

Bases: Interaction