I built cool stuff ai4sci and try to share cool things i saw along the way
@RostislavFedorov
Post #466
92
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
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
- ❤ 3



