search

lab-cosmo / mops

Mathematical operations with GPU acceleration
BSD 3-Clause "New" or "Revised" License
7 stars 1 forks source link

MOPS-Tetris ? #63

Open mitkotak opened 3 months ago

mitkotak commented 3 months ago

Hey e3nn developer jumping in. Thanks for the wonderful work here. Can you point me to some examples or how I might go about implementing a simple nequip tetris model ?

Luthaf commented 3 months ago

Hi! I am not very familiar with the nequip tetris model, but mops works at at different layer than e3nn. Mops provides optimized versions of the very basic mathematical operations, but knows nothing of irreps, equivariants, clebsch-gordan coefficients, …

It would be up to e3nn to use mops internally to implement it's own different layers instead of using native torch operations if you'd want to speed up this tetris model. For example, a layer doing 12x1o \times 8x2e => 10x1o could be implemented using pre-computed CG coefficients and sparse_accumulation_of_products.