feat(jax): freeze to StableXLO & DeepEval

deepmd/backend/jax.py

@@ -34,11 +34,11 @@ class JAXBackend(Backend):
     features: ClassVar[Backend.Feature] = (
         Backend.Feature.IO
         # Backend.Feature.ENTRY_POINT
-        # | Backend.Feature.DEEP_EVAL
+        | Backend.Feature.DEEP_EVAL
         # | Backend.Feature.NEIGHBOR_STAT
     )
     """The features of the backend."""
-    suffixes: ClassVar[list[str]] = [".jax"]
+    suffixes: ClassVar[list[str]] = [".hlo", ".jax"]
     """The suffixes of the backend."""
 
     def is_available(self) -> bool:
@@ -71,7 +71,11 @@ def deep_eval(self) -> type["DeepEvalBackend"]:
         type[DeepEvalBackend]
             The Deep Eval backend of the backend.
         """
-        raise NotImplementedError
+        from deepmd.jax.infer.deep_eval import (
+            DeepEval,
+        )
+
+        return DeepEval
 
     @property
     def neighbor_stat(self) -> type["NeighborStat"]:

deepmd/dpmodel/descriptor/se_e2_a.py

@@ -555,7 +555,7 @@ def call(
             coord_ext, atype_ext, nlist, self.davg, self.dstd
         )
         nf, nloc, nnei, _ = rr.shape
-        sec = xp.asarray(self.sel_cumsum)
+        sec = self.sel_cumsum
 
         ng = self.neuron[-1]
         gr = xp.zeros([nf * nloc, ng, 4], dtype=self.dstd.dtype)

deepmd/dpmodel/utils/serialization.py

@@ -90,7 +90,7 @@ def save_dp_model(filename: str, model_dict: dict) -> None:
         # use UTC+0 time
         "time": str(datetime.datetime.now(tz=datetime.timezone.utc)),
     }
-    if filename_extension == ".dp":
+    if filename_extension in (".dp", ".hlo"):
         variable_counter = Counter()
         with h5py.File(filename, "w") as f:
             model_dict = traverse_model_dict(
@@ -141,7 +141,7 @@ def load_dp_model(filename: str) -> dict:
         The loaded model dict, including meta information.
     """
     filename_extension = Path(filename).suffix
-    if filename_extension == ".dp":
+    if filename_extension in {".dp", ".hlo"}:
         with h5py.File(filename, "r") as f:
             model_dict = json.loads(f.attrs["json"])
             model_dict = traverse_model_dict(model_dict, lambda x: f[x][()].copy())

deepmd/jax/env.py

@@ -8,11 +8,13 @@
 from flax import (
     nnx,
 )
+from jax import export as jax_export
 
 jax.config.update("jax_enable_x64", True)
 
 __all__ = [
     "jax",
     "jnp",
     "nnx",
+    "jax_export",
 ]

deepmd/jax/infer/__init__.py

@@ -0,0 +1 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later