Edge Embedding

We recommend using either IdentityEdgeEmbedding or NonLinearEdgeEmbedding in most cases.

The identity embedding is the most conservative choice. The nonlinear edge embedding is generally recommended when used together with element-dependent edge update.

class tace.models._e3nn.edge.IdentityEdgeEmbedding(num_elements: int, num_radial_basis: int, num_channel: int, bias: bool = False)

Bases: EdgeEmbedding

An identity edge embedding module.

This class directly returns the input edge features (radial) without any transformation.

class tace.models._e3nn.edge.LinearEdgeEmbedding(num_elements: int, num_radial_basis: int, num_channel: int, bias: bool = False)

Bases: EdgeEmbedding

A linear edge embedding module.

This class projects the input edge features (radial) into a higher-dimensional feature space using a linear transformation.

This is motivated by the fact that when edge update are used, a low-dimensional radial representation may become a bottleneck and limit the expressiveness of edge features.

class tace.models._e3nn.edge.NonLinearEdgeEmbedding(num_elements: int, num_radial_basis: int, num_channel: int, bias: bool = False)

Bases: EdgeEmbedding

A nonlinear edge embedding module.

This class applies a nonlinear activation function after a linear projection of edge features, allowing for more expressive representations compared to purely linear transformations.

This is motivated by the fact that when edge update are used, a low-dimensional radial representation may become a bottleneck and limit the expressiveness of edge features.