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.

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

Bases: EdgeEmbedding

An edge embedding module that incorporates both radial and element information.

This class combines transformed edge features with embeddings of the source and target nodes, allowing the edge representation to depend not only on geometric information but also on the types of connected elements.

Note

When using this module, it is recommended not to additionally use edge update modules that rely purely on element information, as this may lead to an overemphasis on element features in the edge representation.