Over the last couple of month i was working on portin e3nn library to MLX arcitecture, so that you can speed up your faivourite flaivour of MLIP on your local Apple Silicon computer.
On my fairly basic MacBook, there is a 5-6× speedups for forward and backward passes using pure MLX, and up to 10× with custom Metal kernels, compared with e3nn on CPU.
To make it easier to get started with the library, we also ported some of the classic e3nn tutorials: the early tutorials by Tess Smidt, as well as the invariants tutorial by Martin Uhrin and Thomas Hardin.
Library:
https://github.com/lamalab-org/e3nn_mlx
Documentation:
https://lamalab-org.github.io/e3nn_mlx/
Tutorials:
https://github.com/lamalab-org/e3nn_mlx/tree/main/tutorials
Post #466
92
- ❤ 3