search

apax-hub / apax

A flexible and performant framework for training machine learning potentials.
https://apax.readthedocs.io
MIT License
18 stars 3 forks source link

cleaned up gm implementation #338

Closed M-R-Schaefer closed 2 months ago

M-R-Schaefer commented 2 months ago

the indexing performed using the TRIL indices was not supported by jax originally (or maybe I was doing something wrong). Whatever the case, it was now possible to significantly clean up the GM implementation. Training results and timings are exactly the same.