feat(jax): checkpoint I/O

deepmd/backend/jax.py

@@ -32,14 +32,13 @@ class JAXBackend(Backend):
     name = "JAX"
     """The formal name of the backend."""
     features: ClassVar[Backend.Feature] = (
-        Backend.Feature(0)
+        Backend.Feature.IO
         # Backend.Feature.ENTRY_POINT
         # | Backend.Feature.DEEP_EVAL
         # | Backend.Feature.NEIGHBOR_STAT
-        # | Backend.Feature.IO
     )
     """The features of the backend."""
-    suffixes: ClassVar[list[str]] = []
+    suffixes: ClassVar[list[str]] = [".jax"]
     """The suffixes of the backend."""
 
     def is_available(self) -> bool:
@@ -94,7 +93,11 @@ def serialize_hook(self) -> Callable[[str], dict]:
         Callable[[str], dict]
             The serialize hook of the backend.
         """
-        raise NotImplementedError
+        from deepmd.jax.utils.serialization import (
+            serialize_from_file,
+        )
+
+        return serialize_from_file
 
     @property
     def deserialize_hook(self) -> Callable[[str, dict], None]:
@@ -105,4 +108,8 @@ def deserialize_hook(self) -> Callable[[str, dict], None]:
         Callable[[str, dict], None]
             The deserialize hook of the backend.
         """
-        raise NotImplementedError
+        from deepmd.jax.utils.serialization import (
+            deserialize_to_file,
+        )
+
+        return deserialize_to_file

deepmd/dpmodel/atomic_model/base_atomic_model.py

@@ -1,13 +1,16 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
 import copy
+import math
 from typing import (
     Optional,
 )
 
+import array_api_compat
 import numpy as np
 
 from deepmd.dpmodel.common import (
     NativeOP,
+    to_numpy_array,
 )
 from deepmd.dpmodel.output_def import (
     FittingOutputDef,
@@ -172,17 +175,18 @@ def forward_common_atomic(
             ret_dict["mask"][ff,ii] == 0 indicating the ii-th atom of the ff-th frame is virtual.
 
         """
+        xp = array_api_compat.array_namespace(extended_coord, extended_atype, nlist)
         _, nloc, _ = nlist.shape
         atype = extended_atype[:, :nloc]
         if self.pair_excl is not None:
             pair_mask = self.pair_excl.build_type_exclude_mask(nlist, extended_atype)
             # exclude neighbors in the nlist
-            nlist = np.where(pair_mask == 1, nlist, -1)
+            nlist = xp.where(pair_mask == 1, nlist, -1)
 
         ext_atom_mask = self.make_atom_mask(extended_atype)
         ret_dict = self.forward_atomic(
             extended_coord,
-            np.where(ext_atom_mask, extended_atype, 0),
+            xp.where(ext_atom_mask, extended_atype, 0),
             nlist,
             mapping=mapping,
             fparam=fparam,
@@ -191,13 +195,13 @@ def forward_common_atomic(
         ret_dict = self.apply_out_stat(ret_dict, atype)
 
         # nf x nloc
-        atom_mask = ext_atom_mask[:, :nloc].astype(np.int32)
+        atom_mask = ext_atom_mask[:, :nloc].astype(xp.int32)
         if self.atom_excl is not None:
             atom_mask *= self.atom_excl.build_type_exclude_mask(atype)
 
         for kk in ret_dict.keys():
             out_shape = ret_dict[kk].shape
-            out_shape2 = np.prod(out_shape[2:])
+            out_shape2 = math.prod(out_shape[2:])
             ret_dict[kk] = (
                 ret_dict[kk].reshape([out_shape[0], out_shape[1], out_shape2])
                 * atom_mask[:, :, None]
@@ -232,8 +236,8 @@ def serialize(self) -> dict:
             "rcond": self.rcond,
             "preset_out_bias": self.preset_out_bias,
             "@variables": {
-                "out_bias": self.out_bias,
-                "out_std": self.out_std,
+                "out_bias": to_numpy_array(self.out_bias),
+                "out_std": to_numpy_array(self.out_std),
             },
         }
 

deepmd/dpmodel/atomic_model/dp_atomic_model.py

@@ -169,14 +169,20 @@ def serialize(self) -> dict:
         )
         return dd
 
+    # for subclass overriden
+    base_descriptor_cls = BaseDescriptor
+    """The base descriptor class."""
+    base_fitting_cls = BaseFitting
+    """The base fitting class."""
+
     @classmethod
     def deserialize(cls, data) -> "DPAtomicModel":
         data = copy.deepcopy(data)
         check_version_compatibility(data.pop("@version", 1), 2, 2)
         data.pop("@class")
         data.pop("type")
-        descriptor_obj = BaseDescriptor.deserialize(data.pop("descriptor"))
-        fitting_obj = BaseFitting.deserialize(data.pop("fitting"))
+        descriptor_obj = cls.base_descriptor_cls.deserialize(data.pop("descriptor"))
+        fitting_obj = cls.base_fitting_cls.deserialize(data.pop("fitting"))
         data["descriptor"] = descriptor_obj
         data["fitting"] = fitting_obj
         obj = super().deserialize(data)

deepmd/dpmodel/model/make_model.py

@@ -3,6 +3,7 @@
     Optional,
 )
 
+import array_api_compat
 import numpy as np
 
 from deepmd.dpmodel.atomic_model.base_atomic_model import (
@@ -75,7 +76,8 @@ def __init__(
             else:
                 self.atomic_model: T_AtomicModel = T_AtomicModel(*args, **kwargs)
             self.precision_dict = PRECISION_DICT
-            self.reverse_precision_dict = RESERVED_PRECISON_DICT
+            # not supported by flax
+            # self.reverse_precision_dict = RESERVED_PRECISON_DICT
             self.global_np_float_precision = GLOBAL_NP_FLOAT_PRECISION
             self.global_ener_float_precision = GLOBAL_ENER_FLOAT_PRECISION
 
@@ -253,9 +255,7 @@ def input_type_cast(
             str,
         ]:
             """Cast the input data to global float type."""
-            input_prec = self.reverse_precision_dict[
-                self.precision_dict[coord.dtype.name]
-            ]
+            input_prec = RESERVED_PRECISON_DICT[self.precision_dict[coord.dtype.name]]
             ###
             ### type checking would not pass jit, convert to coord prec anyway
             ###
@@ -264,10 +264,7 @@ def input_type_cast(
                 for vv in [box, fparam, aparam]
             ]
             box, fparam, aparam = _lst
-            if (
-                input_prec
-                == self.reverse_precision_dict[self.global_np_float_precision]
-            ):
+            if input_prec == RESERVED_PRECISON_DICT[self.global_np_float_precision]:
                 return coord, box, fparam, aparam, input_prec
             else:
                 pp = self.global_np_float_precision
@@ -286,8 +283,7 @@ def output_type_cast(
         ) -> dict[str, np.ndarray]:
             """Convert the model output to the input prec."""
             do_cast = (
-                input_prec
-                != self.reverse_precision_dict[self.global_np_float_precision]
+                input_prec != RESERVED_PRECISON_DICT[self.global_np_float_precision]
             )
             pp = self.precision_dict[input_prec]
             odef = self.model_output_def()
@@ -366,17 +362,18 @@ def _format_nlist(
             nnei: int,
             extra_nlist_sort: bool = False,
         ):
+            xp = array_api_compat.array_namespace(extended_coord, nlist)
             n_nf, n_nloc, n_nnei = nlist.shape
             extended_coord = extended_coord.reshape([n_nf, -1, 3])
             nall = extended_coord.shape[1]
             rcut = self.get_rcut()
 
             if n_nnei < nnei:
                 # make a copy before revise
-                ret = np.concatenate(
+                ret = xp.concat(
                     [
                         nlist,
-                        -1 * np.ones([n_nf, n_nloc, nnei - n_nnei], dtype=nlist.dtype),
+                        -1 * xp.ones([n_nf, n_nloc, nnei - n_nnei], dtype=nlist.dtype),
                     ],
                     axis=-1,
                 )
@@ -385,16 +382,16 @@ def _format_nlist(
                 n_nf, n_nloc, n_nnei = nlist.shape
                 # make a copy before revise
                 m_real_nei = nlist >= 0
-                ret = np.where(m_real_nei, nlist, 0)
+                ret = xp.where(m_real_nei, nlist, 0)
                 coord0 = extended_coord[:, :n_nloc, :]
                 index = ret.reshape(n_nf, n_nloc * n_nnei, 1).repeat(3, axis=2)
-                coord1 = np.take_along_axis(extended_coord, index, axis=1)
+                coord1 = xp.take_along_axis(extended_coord, index, axis=1)
                 coord1 = coord1.reshape(n_nf, n_nloc, n_nnei, 3)
-                rr = np.linalg.norm(coord0[:, :, None, :] - coord1, axis=-1)
-                rr = np.where(m_real_nei, rr, float("inf"))
-                rr, ret_mapping = np.sort(rr, axis=-1), np.argsort(rr, axis=-1)
-                ret = np.take_along_axis(ret, ret_mapping, axis=2)
-                ret = np.where(rr > rcut, -1, ret)
+                rr = xp.linalg.norm(coord0[:, :, None, :] - coord1, axis=-1)
+                rr = xp.where(m_real_nei, rr, float("inf"))
+                rr, ret_mapping = xp.sort(rr, axis=-1), xp.argsort(rr, axis=-1)
+                ret = xp.take_along_axis(ret, ret_mapping, axis=2)
+                ret = xp.where(rr > rcut, -1, ret)
                 ret = ret[..., :nnei]
             # not extra_nlist_sort and n_nnei <= nnei:
             elif n_nnei == nnei:

deepmd/dpmodel/model/transform_output.py

@@ -1,5 +1,6 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
 
+import array_api_compat
 import numpy as np
 
 from deepmd.dpmodel.common import (
@@ -23,6 +24,7 @@ def fit_output_to_model_output(
     the model output.
 
     """
+    xp = array_api_compat.get_namespace(coord_ext)
     model_ret = dict(fit_ret.items())
     for kk, vv in fit_ret.items():
         vdef = fit_output_def[kk]
@@ -31,7 +33,7 @@ def fit_output_to_model_output(
         if vdef.reducible:
             kk_redu = get_reduce_name(kk)
             # cast to energy prec brefore reduction
-            model_ret[kk_redu] = np.sum(
+            model_ret[kk_redu] = xp.sum(
                 vv.astype(GLOBAL_ENER_FLOAT_PRECISION), axis=atom_axis
             )
             if vdef.r_differentiable:

deepmd/jax/atomic_model/__init__.py

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

deepmd/jax/atomic_model/base_atomic_model.py

@@ -0,0 +1,18 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+from deepmd.jax.common import (
+    to_jax_array,
+)
+from deepmd.jax.utils.exclude_mask import (
+    AtomExcludeMask,
+    PairExcludeMask,
+)
+
+
+def base_atomic_model_set_attr(name, value):
+    if name in {"out_bias", "out_std"}:
+        value = to_jax_array(value)
+    elif name == "pair_excl" and value is not None:
+        value = PairExcludeMask(value.ntypes, value.exclude_types)
+    elif name == "atom_excl" and value is not None:
+        value = AtomExcludeMask(value.ntypes, value.exclude_types)
+    return value

deepmd/jax/atomic_model/dp_atomic_model.py

@@ -0,0 +1,30 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+from typing import (
+    Any,
+)
+
+from deepmd.dpmodel.atomic_model.dp_atomic_model import DPAtomicModel as DPAtomicModelDP
+from deepmd.jax.atomic_model.base_atomic_model import (
+    base_atomic_model_set_attr,
+)
+from deepmd.jax.common import (
+    flax_module,
+)
+from deepmd.jax.descriptor.base_descriptor import (
+    BaseDescriptor,
+)
+from deepmd.jax.fitting.base_fitting import (
+    BaseFitting,
+)
+
+
+@flax_module
+class DPAtomicModel(DPAtomicModelDP):
+    base_descriptor_cls = BaseDescriptor
+    """The base descriptor class."""
+    base_fitting_cls = BaseFitting
+    """The base fitting class."""
+
+    def __setattr__(self, name: str, value: Any) -> None:
+        value = base_atomic_model_set_attr(name, value)
+        return super().__setattr__(name, value)

deepmd/jax/descriptor/__init__.py

@@ -1 +1,12 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
+from deepmd.jax.descriptor.dpa1 import (
+    DescrptDPA1,
+)
+from deepmd.jax.descriptor.se_e2_a import (
+    DescrptSeA,
+)
+
+__all__ = [
+    "DescrptSeA",
+    "DescrptDPA1",
+]

deepmd/jax/descriptor/base_descriptor.py

@@ -0,0 +1,9 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+from deepmd.dpmodel.descriptor.make_base_descriptor import (
+    make_base_descriptor,
+)
+from deepmd.jax.env import (
+    jnp,
+)
+
+BaseDescriptor = make_base_descriptor(jnp.ndarray)

deepmd/jax/descriptor/dpa1.py

@@ -16,6 +16,9 @@
     flax_module,
     to_jax_array,
 )
+from deepmd.jax.descriptor.base_descriptor import (
+    BaseDescriptor,
+)
 from deepmd.jax.utils.exclude_mask import (
     PairExcludeMask,
 )
@@ -76,6 +79,8 @@ def __setattr__(self, name: str, value: Any) -> None:
         return super().__setattr__(name, value)
 
 
+@BaseDescriptor.register("dpa1")
+@BaseDescriptor.register("se_atten")
 @flax_module
 class DescrptDPA1(DescrptDPA1DP):
     def __setattr__(self, name: str, value: Any) -> None:

deepmd/jax/descriptor/se_e2_a.py

@@ -8,6 +8,9 @@
     flax_module,
     to_jax_array,
 )
+from deepmd.jax.descriptor.base_descriptor import (
+    BaseDescriptor,
+)
 from deepmd.jax.utils.exclude_mask import (
     PairExcludeMask,
 )
@@ -16,6 +19,8 @@
 )
 
 
+@BaseDescriptor.register("se_e2_a")
+@BaseDescriptor.register("se_a")
 @flax_module
 class DescrptSeA(DescrptSeADP):
     def __setattr__(self, name: str, value: Any) -> None:

deepmd/jax/env.py

@@ -8,11 +8,15 @@
 from flax import (
     nnx,
 )
+from jax.experimental import (
+    jax2tf,
+)
 
 jax.config.update("jax_enable_x64", True)
 
 __all__ = [
     "jax",
     "jnp",
     "nnx",
+    "jax2tf",
 ]

deepmd/jax/fitting/__init__.py

@@ -1 +1,10 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
+from deepmd.jax.fitting.fitting import (
+    DOSFittingNet,
+    EnergyFittingNet,
+)
+
+__all__ = [
+    "EnergyFittingNet",
+    "DOSFittingNet",
+]

deepmd/jax/fitting/base_fitting.py

@@ -0,0 +1,9 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+from deepmd.dpmodel.fitting.make_base_fitting import (
+    make_base_fitting,
+)
+from deepmd.jax.env import (
+    jnp,
+)
+
+BaseFitting = make_base_fitting(jnp.ndarray)