jax-dimenet (pypi): DimeNet++ in Jax.
Keywords: graph-neural-networks|jax|molecular-dynamics