gatr_linear

cuequivariance.descriptors.gatr_linear(

mul_in: int,

mul_out: int,

) → SegmentedTensorProduct

subsrcipts: weights[uv],input[iu],output[iv]

References

• https://arxiv.org/pdf/2305.18415

• Source code

Examples

>>> d = gatr_linear(32, 32)
>>> d
uv,iu,iv sizes=9216,512,512 num_segments=9,8,8 num_paths=12 i={1, 3} u=32 v=32