[Feature Request] Refactor the cpp interface of the saved DMFF jax model with MD engine

[dingye18]

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.

1. 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
2. Unsupported data type f64, s64,

Suggested Solutions

https://github.com/google/jax/issues/1871 Old solution.

Lack of documentation, more exploration required

Further Information, Files, and Links

No response