e3x.nn.modules

Equivariant neural network modules.

Classes

Dense(features[, use_bias, dtype, ...])	A linear transformation applied over the last dimension of the input.
Embed(num_embeddings, features[, dtype, ...])	Embedding module.
FusedTensor([max_degree, ...])	Fused tensor product of two equivariant feature representations.
MessagePass([max_degree, use_basis_bias, ...])	Equivariant message-passing step.
MultiHeadAttention([max_degree, ...])	Equivariant multi-head attention.
SelfAttention([max_degree, use_basis_bias, ...])	Equivariant self-attention.
Tensor([max_degree, include_pseudotensors, ...])	Tensor product of two equivariant feature representations.
TensorDense([features, max_degree, ...])	Linear projection followed by a tensor product.