fix(jax): fix compatibility with flax 0.12

Author: njzjz

.github/workflows/test_python.yml

@@ -32,7 +32,7 @@ jobs:
         source/install/uv_with_retry.sh pip install --system torch -i https://download.pytorch.org/whl/cpu
         export TENSORFLOW_ROOT=$(python -c 'import importlib.util,pathlib;print(pathlib.Path(importlib.util.find_spec("tensorflow").origin).parent)')
         export PYTORCH_ROOT=$(python -c 'import torch;print(torch.__path__[0])')
-        source/install/uv_with_retry.sh pip install --system -e .[test,jax] mpi4py "jax==0.5.0;python_version>='3.10'"
+        source/install/uv_with_retry.sh pip install --system -e .[test,jax] mpi4py "jax==0.8.1;python_version>='3.10'"
         source/install/uv_with_retry.sh pip install --system -U setuptools
         source/install/uv_with_retry.sh pip install --system horovod --no-build-isolation
         source/install/uv_with_retry.sh pip install --system --pre "paddlepaddle==3.0.0" -i https://www.paddlepaddle.org.cn/packages/stable/cpu/

deepmd/jax/atomic_model/linear_atomic_model.py

@@ -25,6 +25,13 @@
     jax,
     jnp,
 )
+from packaging.version import (
+    Version,
+)
+from deepmd.jax.env import (
+    flax_version,
+    nnx,
+)
 
 
 @flax_module
@@ -33,13 +40,17 @@ def __setattr__(self, name: str, value: Any) -> None:
         value = base_atomic_model_set_attr(name, value)
         if name == "mapping_list":
             value = [ArrayAPIVariable(to_jax_array(vv)) for vv in value]
+            if Version(flax_version) >= Version("0.12.0"):
+                value = nnx.List([nnx.data(item) for item in value])
         elif name == "zbl_weight":
             value = ArrayAPIVariable(to_jax_array(value))
         elif name == "models":
             value = [
                 DPAtomicModel.deserialize(value[0].serialize()),
                 PairTabAtomicModel.deserialize(value[1].serialize()),
             ]
+            if Version(flax_version) >= Version("0.12.0"):
+                value = nnx.List([nnx.data(item) for item in value])
         return super().__setattr__(name, value)
 
     def forward_common_atomic(

deepmd/jax/descriptor/dpa1.py

@@ -31,6 +31,13 @@
 from deepmd.jax.utils.type_embed import (
     TypeEmbedNet,
 )
+from packaging.version import (
+    Version,
+)
+from deepmd.jax.env import (
+    flax_version,
+    nnx,
+)
 
 
 @flax_module
@@ -58,6 +65,8 @@ def __setattr__(self, name: str, value: Any) -> None:
             value = [
                 NeighborGatedAttentionLayer.deserialize(ii.serialize()) for ii in value
             ]
+            if Version(flax_version) >= Version("0.12.0"):
+                value = nnx.List([nnx.data(item) for item in value])
         return super().__setattr__(name, value)
 
 
@@ -71,6 +80,8 @@ def __setattr__(self, name: str, value: Any) -> None:
         elif name in {"embeddings", "embeddings_strip"}:
             if value is not None:
                 value = NetworkCollection.deserialize(value.serialize())
+            elif Version(flax_version) >= Version("0.12.0"):
+                value = nnx.data(value)
         elif name == "dpa1_attention":
             value = NeighborGatedAttention.deserialize(value.serialize())
         elif name == "env_mat":

deepmd/jax/descriptor/dpa2.py

@@ -29,6 +29,13 @@
 from deepmd.jax.utils.type_embed import (
     TypeEmbedNet,
 )
+from packaging.version import (
+    Version,
+)
+from deepmd.jax.env import (
+    flax_version,
+    nnx,
+)
 
 
 @BaseDescriptor.register("dpa2")
@@ -44,13 +51,16 @@ def __setattr__(self, name: str, value: Any) -> None:
         elif name in {"repinit_three_body"}:
             if value is not None:
                 value = DescrptBlockSeTTebd.deserialize(value.serialize())
+            elif Version(flax_version) >= Version("0.12.0"):
+                value = nnx.data(value)
         elif name in {"repformers"}:
             value = DescrptBlockRepformers.deserialize(value.serialize())
         elif name in {"type_embedding"}:
             value = TypeEmbedNet.deserialize(value.serialize())
         elif name in {"g1_shape_tranform", "tebd_transform"}:
             if value is None:
-                pass
+                if Version(flax_version) >= Version("0.12.0"):
+                    value = nnx.data(value)
             elif isinstance(value, NativeLayerDP):
                 value = NativeLayer.deserialize(value.serialize())
             elif isinstance(value, IdentityDP):

deepmd/jax/descriptor/hybrid.py

@@ -12,15 +12,25 @@
 from deepmd.jax.descriptor.base_descriptor import (
     BaseDescriptor,
 )
-
+from packaging.version import (
+    Version,
+)
+from deepmd.jax.env import (
+    flax_version,
+    nnx,
+)
 
 @BaseDescriptor.register("hybrid")
 @flax_module
 class DescrptHybrid(DescrptHybridDP):
     def __setattr__(self, name: str, value: Any) -> None:
         if name in {"nlist_cut_idx"}:
             value = [ArrayAPIVariable(to_jax_array(vv)) for vv in value]
+            if Version(flax_version) >= Version("0.12.0"):
+                value = nnx.List([nnx.data(item) for item in value])
         elif name in {"descrpt_list"}:
             value = [BaseDescriptor.deserialize(vv.serialize()) for vv in value]
+            if Version(flax_version) >= Version("0.12.0"):
+                value = nnx.List([nnx.data(item) for item in value])
 
         return super().__setattr__(name, value)

deepmd/jax/descriptor/repflows.py

@@ -18,6 +18,13 @@
 from deepmd.jax.utils.network import (
     NativeLayer,
 )
+from packaging.version import (
+    Version,
+)
+from deepmd.jax.env import (
+    flax_version,
+    nnx,
+)
 
 
 @flax_module
@@ -29,6 +36,8 @@ def __setattr__(self, name: str, value: Any) -> None:
                 value = ArrayAPIVariable(value)
         elif name in {"layers"}:
             value = [RepFlowLayer.deserialize(layer.serialize()) for layer in value]
+            if Version(flax_version) >= Version("0.12.0"):
+                value = nnx.List([nnx.data(item) for item in value])
         elif name in {"edge_embd", "angle_embd"}:
             value = NativeLayer.deserialize(value.serialize())
         elif name in {"env_mat_edge", "env_mat_angle"}:
@@ -58,8 +67,12 @@ def __setattr__(self, name: str, value: Any) -> None:
         }:
             if value is not None:
                 value = NativeLayer.deserialize(value.serialize())
+            elif Version(flax_version) >= Version("0.12.0"):
+                value = nnx.data(value)
         elif name in {"n_residual", "e_residual", "a_residual"}:
             value = [ArrayAPIVariable(to_jax_array(vv)) for vv in value]
+            if Version(flax_version) >= Version("0.12.0"):
+                value = nnx.List([nnx.data(item) for item in value])
         else:
             pass
         return super().__setattr__(name, value)

deepmd/jax/descriptor/repformers.py

@@ -27,6 +27,13 @@
     LayerNorm,
     NativeLayer,
 )
+from packaging.version import (
+    Version,
+)
+from deepmd.jax.env import (
+    flax_version,
+    nnx,
+)
 
 
 @flax_module
@@ -38,6 +45,8 @@ def __setattr__(self, name: str, value: Any) -> None:
                 value = ArrayAPIVariable(value)
         elif name in {"layers"}:
             value = [RepformerLayer.deserialize(layer.serialize()) for layer in value]
+            if Version(flax_version) >= Version("0.12.0"):
+                value = nnx.List([nnx.data(item) for item in value])
         elif name == "g2_embd":
             value = NativeLayer.deserialize(value.serialize())
         elif name == "env_mat":
@@ -87,21 +96,35 @@ def __setattr__(self, name: str, value: Any) -> None:
         if name in {"linear1", "linear2", "g1_self_mlp", "proj_g1g2", "proj_g1g1g2"}:
             if value is not None:
                 value = NativeLayer.deserialize(value.serialize())
+            elif Version(flax_version) >= Version("0.12.0"):
+                value = nnx.data(value)
         elif name in {"g1_residual", "g2_residual", "h2_residual"}:
             value = [ArrayAPIVariable(to_jax_array(vv)) for vv in value]
+            if Version(flax_version) >= Version("0.12.0"):
+                value = nnx.List([nnx.data(item) for item in value])
         elif name in {"attn2g_map"}:
             if value is not None:
                 value = Atten2Map.deserialize(value.serialize())
+            elif Version(flax_version) >= Version("0.12.0"):
+                value = nnx.data(value)
         elif name in {"attn2_mh_apply"}:
             if value is not None:
                 value = Atten2MultiHeadApply.deserialize(value.serialize())
+            elif Version(flax_version) >= Version("0.12.0"):
+                value = nnx.data(value)
         elif name in {"attn2_lm"}:
             if value is not None:
                 value = LayerNorm.deserialize(value.serialize())
+            elif Version(flax_version) >= Version("0.12.0"):
+                value = nnx.data(value)
         elif name in {"attn2_ev_apply"}:
             if value is not None:
                 value = Atten2EquiVarApply.deserialize(value.serialize())
+            elif Version(flax_version) >= Version("0.12.0"):
+                value = nnx.data(value)
         elif name in {"loc_attn"}:
             if value is not None:
                 value = LocalAtten.deserialize(value.serialize())
+            elif Version(flax_version) >= Version("0.12.0"):
+                value = nnx.data(value)
         return super().__setattr__(name, value)

deepmd/jax/descriptor/se_e2_a.py

@@ -19,6 +19,14 @@
     NetworkCollection,
 )
 
+from packaging.version import (
+    Version,
+)
+from deepmd.jax.env import (
+    flax_version,
+    nnx,
+)
+
 
 @BaseDescriptor.register("se_e2_a")
 @BaseDescriptor.register("se_a")
@@ -32,6 +40,8 @@ def __setattr__(self, name: str, value: Any) -> None:
         elif name in {"embeddings"}:
             if value is not None:
                 value = NetworkCollection.deserialize(value.serialize())
+            elif Version(flax_version) >= Version("0.12.0"):
+                value = nnx.data(value)
         elif name == "env_mat":
             # env_mat doesn't store any value
             pass

deepmd/jax/descriptor/se_e2_r.py

@@ -19,6 +19,14 @@
     NetworkCollection,
 )
 
+from packaging.version import (
+    Version,
+)
+from deepmd.jax.env import (
+    flax_version,   
+    nnx,
+)
+
 
 @BaseDescriptor.register("se_e2_r")
 @BaseDescriptor.register("se_r")
@@ -32,6 +40,8 @@ def __setattr__(self, name: str, value: Any) -> None:
         elif name in {"embeddings"}:
             if value is not None:
                 value = NetworkCollection.deserialize(value.serialize())
+            elif Version(flax_version) >= Version("0.12.0"):
+                value = nnx.data(value)
         elif name == "env_mat":
             # env_mat doesn't store any value
             pass

deepmd/jax/env.py

@@ -7,6 +7,7 @@
 import jax.numpy as jnp
 from flax import (
     nnx,
+    __version__ as flax_version,
 )
 from jax import export as jax_export
 
@@ -23,4 +24,5 @@
     "jax_export",
     "jnp",
     "nnx",
+    "flax_version",
 ]