Create neural networks equivariant to reductive Lie groups in Pytorch. This library implements the Pytorch functions to use all the group operations implemented in lie-nn on torch tensors data.
Tensor products between two irreducible representations of Lie group, reduced into a sum of irreducible representations.
irrep1 = [lie.MulIrrep(3, lie.irreps.SU3((2, 1, 0))), lie.MulIrrep(3, lie.irreps.SU3((1, 1, 0)))]
irrep2 = [lie.MulIrrep(1, lie.irreps.SU3((1, 1, 0)))]
irrep3 = [lie.MulIrrep(1, lie.irreps.SU3((1, 0, 0)))]
irreps1 = lie.ReducedRep.from_irreps(irrep1)
irreps2 = lie.ReducedRep.from_irreps(irrep2)
irreps3 = lie.ReducedRep.from_irreps(irrep3)
irrep_out, instructions = tp_out_irreps_with_instructions(irreps1.irreps, irreps2.irreps, irreps3.irreps)
tp = TensorProduct(irreps1.irreps, irreps2.irreps, irrep_out.irreps, instructions)
x1 = torch.randn(1, irreps1.dim)
x2 = torch.randn(1, irreps2.dim)
out = tp(x1, x2)
irrep1 = [lie.MulIrrep(3, lie.irreps.SU3((2, 1, 0))), lie.MulIrrep(3, lie.irreps.SU3((1, 1, 0)))]
irrep2 = [lie.MulIrrep(5, lie.irreps.SU3((1, 1, 0)))]
irreps1 = lie.ReducedRep.from_irreps(irrep1)
irreps2 = lie.ReducedRep.from_irreps(irrep2)
linear = Linear(irreps1, irreps2)
x1 = torch.randn(1, irreps1.dim)
out = linear(x1)
pycln .
black .
If you use this code, please cite our papers:
@misc{batatia2023general,
title={A General Framework for Equivariant Neural Networks on Reductive Lie Groups},
author={Ilyes Batatia and Mario Geiger and Jose Munoz and Tess Smidt and Lior Silberman and Christoph Ortner},
year={2023},
eprint={2306.00091},
archivePrefix={arXiv},
primaryClass={stat.ML}
}