Reduction of Tensors in Irreps

class e3nn.o3.ReducedTensorProducts(formula, filter_ir_out=None, filter_ir_mid=None, eps: float = 1e-09, **irreps)

Bases: CodeGenMixin, Module

reduce a tensor with symmetries into irreducible representations

Parameters:

• formula (str) – String made of letters - and = that represent the indices symmetries of the tensor. For instance ij=ji means that the tensor has two indices and if they are exchanged, its value is the same. ij=-ji means that the tensor change its sign if the two indices are exchanged.

• filter_ir_out (list of e3nn.o3.Irrep, optional) – Optional, list of allowed irrep in the output

• filter_ir_mid (list of e3nn.o3.Irrep, optional) – Optional, list of allowed irrep in the intermediary operations

• **kwargs (dict of e3nn.o3.Irreps) – each letter present in the formula has to be present in the irreps dictionary, unless it can be inferred by the formula. For instance if the formula is ij=ji you can provide the representation of i only: ReducedTensorProducts('ij=ji', i='1o').

irreps_in

input representations

Type:

list of e3nn.o3.Irreps

irreps_out

output representation

Type:

e3nn.o3.Irreps

change_of_basis

tensor of shape (irreps_out.dim, irreps_in[0].dim, ..., irreps_in[-1].dim)

Type:

torch.Tensor

Examples

>>> tp = ReducedTensorProducts('ij=-ji', i='1o')
>>> x = torch.tensor([1.0, 0.0, 0.0])
>>> y = torch.tensor([0.0, 1.0, 0.0])
>>> tp(x, y) + tp(y, x)
tensor([0., 0., 0.])

>>> tp = ReducedTensorProducts('ijkl=jikl=ikjl=ijlk', i="1e")
>>> tp.irreps_out
1x0e+1x2e+1x4e

>>> tp = ReducedTensorProducts('ij=ji', i='1o')
>>> x, y = torch.randn(2, 3)
>>> a = torch.einsum('zij,i,j->z', tp.change_of_basis, x, y)
>>> b = tp(x, y)
>>> assert torch.allclose(a, b, atol=1e-3, rtol=1e-3)

Methods:

forward(*xs)

Define the computation performed at every call.

forward(*xs)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.