Summary
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.
- Limited support for custom calls. https://github.com/google/jax/tree/jaxlib-v0.4.25/jax/experimental/jax2tf#native-serialization-supports-only-select-custom-calls.
Occurred when using JAX 0.4.24 + TF 2.15/2.14
- Unsupported data type f64, s64,
Suggested Solutions
jax-ml/jax#1871
Old solution.
Lack of documentation, more exploration required
Further Information, Files, and Links
No response
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.Occurred when using JAX 0.4.24 + TF 2.15/2.14
Suggested Solutions
jax-ml/jax#1871
Old solution.
Lack of documentation, more exploration required
Further Information, Files, and Links
No response