feat(jax): checkpoint I/O (#4236)

Implement a JAX checkpoint format. I name it `*.jax` as I don't find
existing conventions.

<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit

## Release Notes

- **New Features**
- Introduced serialization and deserialization functionalities for JAX
backend models.
- Added support for the `.jax` file suffix in the backend configuration.
- Enhanced attribute handling logic across various classes to ensure
proper processing of non-null values.

- **Bug Fixes**
	- Enhanced cleanup processes in the test suite to improve reliability.

- **Chores**
- Updated dependencies in the project configuration for better JAX
compatibility.
	- Adjusted linting rules to accommodate JAX-related code.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Signed-off-by: Jinzhe Zeng <[email protected]>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: njzjz

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/jax/atomic_model/base_atomic_model.py

@@ -1,5 +1,6 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
 from deepmd.jax.common import (
+    ArrayAPIVariable,
     to_jax_array,
 )
 from deepmd.jax.utils.exclude_mask import (
@@ -11,6 +12,8 @@
 def base_atomic_model_set_attr(name, value):
     if name in {"out_bias", "out_std"}:
         value = to_jax_array(value)
+        if value is not None:
+            value = ArrayAPIVariable(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:

deepmd/jax/common.py

@@ -81,3 +81,17 @@ def __setattr__(self, name: str, value: Any) -> None:
             return super().__setattr__(name, value)
 
     return FlaxModule
+
+
+class ArrayAPIVariable(nnx.Variable):
+    def __array__(self, *args, **kwargs):
+        return self.value.__array__(*args, **kwargs)
+
+    def __array_namespace__(self, *args, **kwargs):
+        return self.value.__array_namespace__(*args, **kwargs)
+
+    def __dlpack__(self, *args, **kwargs):
+        return self.value.__dlpack__(*args, **kwargs)
+
+    def __dlpack_device__(self, *args, **kwargs):
+        return self.value.__dlpack_device__(*args, **kwargs)

deepmd/jax/descriptor/dpa1.py

@@ -13,6 +13,7 @@
     NeighborGatedAttentionLayer as NeighborGatedAttentionLayerDP,
 )
 from deepmd.jax.common import (
+    ArrayAPIVariable,
     flax_module,
     to_jax_array,
 )
@@ -65,6 +66,8 @@ class DescrptBlockSeAtten(DescrptBlockSeAttenDP):
     def __setattr__(self, name: str, value: Any) -> None:
         if name in {"mean", "stddev"}:
             value = to_jax_array(value)
+            if value is not None:
+                value = ArrayAPIVariable(value)
         elif name in {"embeddings", "embeddings_strip"}:
             if value is not None:
                 value = NetworkCollection.deserialize(value.serialize())

deepmd/jax/descriptor/se_e2_a.py

@@ -5,6 +5,7 @@
 
 from deepmd.dpmodel.descriptor.se_e2_a import DescrptSeAArrayAPI as DescrptSeADP
 from deepmd.jax.common import (
+    ArrayAPIVariable,
     flax_module,
     to_jax_array,
 )
@@ -26,6 +27,8 @@ class DescrptSeA(DescrptSeADP):
     def __setattr__(self, name: str, value: Any) -> None:
         if name in {"dstd", "davg"}:
             value = to_jax_array(value)
+            if value is not None:
+                value = ArrayAPIVariable(value)
         elif name in {"embeddings"}:
             if value is not None:
                 value = NetworkCollection.deserialize(value.serialize())

deepmd/jax/fitting/fitting.py

@@ -6,6 +6,7 @@
 from deepmd.dpmodel.fitting.dos_fitting import DOSFittingNet as DOSFittingNetDP
 from deepmd.dpmodel.fitting.ener_fitting import EnergyFittingNet as EnergyFittingNetDP
 from deepmd.jax.common import (
+    ArrayAPIVariable,
     flax_module,
     to_jax_array,
 )
@@ -29,6 +30,8 @@ def setattr_for_general_fitting(name: str, value: Any) -> Any:
         "aparam_inv_std",
     }:
         value = to_jax_array(value)
+        if value is not None:
+            value = ArrayAPIVariable(value)
     elif name == "emask":
         value = AtomExcludeMask(value.ntypes, value.exclude_types)
     elif name == "nets":

deepmd/jax/utils/exclude_mask.py

@@ -6,6 +6,7 @@
 from deepmd.dpmodel.utils.exclude_mask import AtomExcludeMask as AtomExcludeMaskDP
 from deepmd.dpmodel.utils.exclude_mask import PairExcludeMask as PairExcludeMaskDP
 from deepmd.jax.common import (
+    ArrayAPIVariable,
     flax_module,
     to_jax_array,
 )
@@ -16,6 +17,8 @@ class AtomExcludeMask(AtomExcludeMaskDP):
     def __setattr__(self, name: str, value: Any) -> None:
         if name in {"type_mask"}:
             value = to_jax_array(value)
+            if value is not None:
+                value = ArrayAPIVariable(value)
         return super().__setattr__(name, value)
 
 
@@ -24,4 +27,6 @@ class PairExcludeMask(PairExcludeMaskDP):
     def __setattr__(self, name: str, value: Any) -> None:
         if name in {"type_mask"}:
             value = to_jax_array(value)
+            if value is not None:
+                value = ArrayAPIVariable(value)
         return super().__setattr__(name, value)

deepmd/jax/utils/serialization.py

@@ -0,0 +1,97 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+from pathlib import (
+    Path,
+)
+
+import orbax.checkpoint as ocp
+
+from deepmd.jax.env import (
+    jax,
+    nnx,
+)
+from deepmd.jax.model.model import (
+    BaseModel,
+    get_model,
+)
+
+
+def deserialize_to_file(model_file: str, data: dict) -> None:
+    """Deserialize the dictionary to a model file.
+
+    Parameters
+    ----------
+    model_file : str
+        The model file to be saved.
+    data : dict
+        The dictionary to be deserialized.
+    """
+    if model_file.endswith(".jax"):
+        model = BaseModel.deserialize(data["model"])
+        model_def_script = data["model_def_script"]
+        _, state = nnx.split(model)
+        with ocp.Checkpointer(
+            ocp.CompositeCheckpointHandler("state", "model_def_script")
+        ) as checkpointer:
+            checkpointer.save(
+                Path(model_file).absolute(),
+                ocp.args.Composite(
+                    state=ocp.args.StandardSave(state.to_pure_dict()),
+                    model_def_script=ocp.args.JsonSave(model_def_script),
+                ),
+            )
+    else:
+        raise ValueError("JAX backend only supports converting .jax directory")
+
+
+def serialize_from_file(model_file: str) -> dict:
+    """Serialize the model file to a dictionary.
+
+    Parameters
+    ----------
+    model_file : str
+        The model file to be serialized.
+
+    Returns
+    -------
+    dict
+        The serialized model data.
+    """
+    if model_file.endswith(".jax"):
+        with ocp.Checkpointer(
+            ocp.CompositeCheckpointHandler("state", "model_def_script")
+        ) as checkpointer:
+            data = checkpointer.restore(
+                Path(model_file).absolute(),
+                ocp.args.Composite(
+                    state=ocp.args.StandardRestore(),
+                    model_def_script=ocp.args.JsonRestore(),
+                ),
+            )
+        state = data.state
+
+        # convert str "1" to int 1 key
+        def convert_str_to_int_key(item: dict):
+            for key, value in item.copy().items():
+                if isinstance(value, dict):
+                    convert_str_to_int_key(value)
+                if key.isdigit():
+                    item[int(key)] = item.pop(key)
+
+        convert_str_to_int_key(state)
+
+        model_def_script = data.model_def_script
+        abstract_model = get_model(model_def_script)
+        graphdef, abstract_state = nnx.split(abstract_model)
+        abstract_state.replace_by_pure_dict(state)
+        model = nnx.merge(graphdef, abstract_state)
+        model_dict = model.serialize()
+        data = {
+            "backend": "JAX",
+            "jax_version": jax.__version__,
+            "model": model_dict,
+            "model_def_script": model_def_script,
+            "@variables": {},
+        }
+        return data
+    else:
+        raise ValueError("JAX backend only supports converting .jax directory")

deepmd/jax/utils/type_embed.py

@@ -5,6 +5,7 @@
 
 from deepmd.dpmodel.utils.type_embed import TypeEmbedNet as TypeEmbedNetDP
 from deepmd.jax.common import (
+    ArrayAPIVariable,
     flax_module,
     to_jax_array,
 )
@@ -18,6 +19,8 @@ class TypeEmbedNet(TypeEmbedNetDP):
     def __setattr__(self, name: str, value: Any) -> None:
         if name in {"econf_tebd"}:
             value = to_jax_array(value)
+            if value is not None:
+                value = ArrayAPIVariable(value)
         if name in {"embedding_net"}:
             value = EmbeddingNet.deserialize(value.serialize())
         return super().__setattr__(name, value)

pyproject.toml

@@ -137,7 +137,10 @@ cu12 = [
 ]
 jax = [
     'jax>=0.4.33;python_version>="3.10"',
-    'flax>=0.8.0;python_version>="3.10"',
+    'flax>=0.10.0;python_version>="3.10"',
+    'orbax-checkpoint;python_version>="3.10"',
+    # The pinning of ml_dtypes may conflict with TF
+    # 'jax-ai-stack;python_version>="3.10"',
 ]
 
 [tool.deepmd_build_backend.scripts]
@@ -402,6 +405,7 @@ banned-module-level-imports = [
 # Also ignore `E402` in all `__init__.py` files.
 "deepmd/tf/**" = ["TID253"]
 "deepmd/pt/**" = ["TID253"]
+"deepmd/jax/**" = ["TID253"]
 "source/tests/tf/**" = ["TID253"]
 "source/tests/pt/**" = ["TID253"]
 "source/tests/universal/pt/**" = ["TID253"]