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

### 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](https://github.com/google/jax/tree/jaxlib-v0.4.25/jax/experimental/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_