In case you missed it, here is a similar implementation https://github.com/ACEsuit/mace-jax/tree/main
