DMFF (Differentiable Molecular Force Field) is a Jax-based python package that provides a full differentiable implementation of molecular force field models.
GNU Lesser General Public License v3.0
154
stars
43
forks
source link
[Feature Request] Refactor the cpp interface of the saved DMFF jax model with MD engine #173
Moving the jax2tf to HLOModule for the cpp interface of the saved DMFF model.
Motivation
The current implementation of the cpp interface between the saved DMFF model and MD engine was based on the jax2tf.
The jax2tf was used to convert the the jax function to TensorFlow function.
However, as an experimental feature of JAX, jax2tf does have some limitations for production use.
Summary
Moving the
jax2tftoHLOModulefor the cpp interface of the saved DMFF model.Motivation
The current implementation of the cpp interface between the saved DMFF model and MD engine was based on the jax2tf. The
jax2tfwas used to convert the the jax function to TensorFlow function. However, as an experimental feature of JAX,jax2tfdoes have some limitations for production use.Suggested Solutions
https://github.com/google/jax/issues/1871 Old solution.
Lack of documentation, more exploration required
Further Information, Files, and Links
No response