fix(jax): fix compatibility with flax 0.12 (#5067)

Since the latest TF and JAX have not been compatible with each other, I
keep the old JAX and flax version in the CI.

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

## Summary by CodeRabbit

## Release Notes

* **Chores**
* Updated internal handling across model and descriptor modules to
support Flax 0.12.0 and later versions with conditional runtime behavior
based on detected Flax version.

<sub>✏️ Tip: You can customize this high-level summary in your review
settings.</sub>

<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Signed-off-by: Jinzhe Zeng <jinzhe.zeng@ustc.edu.cn>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: njzjz

deepmd/jax/atomic_model/base_atomic_model.py

@@ -3,10 +3,18 @@
     Any,
 )
 
+from packaging.version import (
+    Version,
+)
+
 from deepmd.jax.common import (
     ArrayAPIVariable,
     to_jax_array,
 )
+from deepmd.jax.env import (
+    flax_version,
+    nnx,
+)
 from deepmd.jax.utils.exclude_mask import (
     AtomExcludeMask,
     PairExcludeMask,
@@ -18,6 +26,8 @@ def base_atomic_model_set_attr(name: str, value: Any) -> Any:
         value = to_jax_array(value)
         if value is not None:
             value = ArrayAPIVariable(value)
+        elif Version(flax_version) >= Version("0.12.0"):
+            value = nnx.data(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/atomic_model/linear_atomic_model.py

@@ -4,6 +4,10 @@
     Optional,
 )
 
+from packaging.version import (
+    Version,
+)
+
 from deepmd.dpmodel.atomic_model.linear_atomic_model import (
     DPZBLLinearEnergyAtomicModel as DPZBLLinearEnergyAtomicModelDP,
 )
@@ -22,8 +26,10 @@
     to_jax_array,
 )
 from deepmd.jax.env import (
+    flax_version,
     jax,
     jnp,
+    nnx,
 )
 
 
@@ -33,13 +39,19 @@ 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))
+            # discard since it's only used in tests
+            # to fix flax.errors.TraceContextError: Cannot mutate 'FlaxModule' from different trace level
+            return
         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/atomic_model/pairtab_atomic_model.py

@@ -4,6 +4,10 @@
     Optional,
 )
 
+from packaging.version import (
+    Version,
+)
+
 from deepmd.dpmodel.atomic_model.pairtab_atomic_model import (
     PairTabAtomicModel as PairTabAtomicModelDP,
 )
@@ -16,8 +20,10 @@
     to_jax_array,
 )
 from deepmd.jax.env import (
+    flax_version,
     jax,
     jnp,
+    nnx,
 )
 
 
@@ -29,6 +35,8 @@ def __setattr__(self, name: str, value: Any) -> None:
             value = to_jax_array(value)
             if value is not None:
                 value = ArrayAPIVariable(value)
+            elif Version(flax_version) >= Version("0.12.0"):
+                value = nnx.data(value)
         return super().__setattr__(name, value)
 
     def forward_common_atomic(

deepmd/jax/descriptor/dpa1.py

@@ -3,6 +3,10 @@
     Any,
 )
 
+from packaging.version import (
+    Version,
+)
+
 from deepmd.dpmodel.descriptor.dpa1 import DescrptBlockSeAtten as DescrptBlockSeAttenDP
 from deepmd.dpmodel.descriptor.dpa1 import DescrptDPA1 as DescrptDPA1DP
 from deepmd.dpmodel.descriptor.dpa1 import GatedAttentionLayer as GatedAttentionLayerDP
@@ -20,6 +24,10 @@
 from deepmd.jax.descriptor.base_descriptor import (
     BaseDescriptor,
 )
+from deepmd.jax.env import (
+    flax_version,
+    nnx,
+)
 from deepmd.jax.utils.exclude_mask import (
     PairExcludeMask,
 )
@@ -58,6 +66,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)
 
 
@@ -68,9 +78,13 @@ def __setattr__(self, name: str, value: Any) -> None:
             value = to_jax_array(value)
             if value is not None:
                 value = ArrayAPIVariable(value)
+            elif Version(flax_version) >= Version("0.12.0"):
+                value = nnx.data(value)
         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

@@ -3,6 +3,10 @@
     Any,
 )
 
+from packaging.version import (
+    Version,
+)
+
 from deepmd.dpmodel.descriptor.dpa2 import DescrptDPA2 as DescrptDPA2DP
 from deepmd.dpmodel.utils.network import Identity as IdentityDP
 from deepmd.dpmodel.utils.network import NativeLayer as NativeLayerDP
@@ -23,6 +27,10 @@
 from deepmd.jax.descriptor.se_t_tebd import (
     DescrptBlockSeTTebd,
 )
+from deepmd.jax.env import (
+    flax_version,
+    nnx,
+)
 from deepmd.jax.utils.network import (
     NativeLayer,
 )
@@ -39,18 +47,23 @@ def __setattr__(self, name: str, value: Any) -> None:
             value = to_jax_array(value)
             if value is not None:
                 value = ArrayAPIVariable(value)
+            elif Version(flax_version) >= Version("0.12.0"):
+                value = nnx.data(value)
         elif name in {"repinit"}:
             value = DescrptBlockSeAtten.deserialize(value.serialize())
         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/dpa3.py

@@ -3,6 +3,10 @@
     Any,
 )
 
+from packaging.version import (
+    Version,
+)
+
 from deepmd.dpmodel.descriptor.dpa3 import DescrptDPA3 as DescrptDPA3DP
 from deepmd.jax.common import (
     ArrayAPIVariable,
@@ -15,6 +19,10 @@
 from deepmd.jax.descriptor.repflows import (
     DescrptBlockRepflows,
 )
+from deepmd.jax.env import (
+    flax_version,
+    nnx,
+)
 from deepmd.jax.utils.type_embed import (
     TypeEmbedNet,
 )
@@ -28,6 +36,8 @@ def __setattr__(self, name: str, value: Any) -> None:
             value = to_jax_array(value)
             if value is not None:
                 value = ArrayAPIVariable(value)
+            elif Version(flax_version) >= Version("0.12.0"):
+                value = nnx.data(value)
         elif name in {"repflows"}:
             value = DescrptBlockRepflows.deserialize(value.serialize())
         elif name in {"type_embedding"}:

deepmd/jax/descriptor/hybrid.py

@@ -3,6 +3,10 @@
     Any,
 )
 
+from packaging.version import (
+    Version,
+)
+
 from deepmd.dpmodel.descriptor.hybrid import DescrptHybrid as DescrptHybridDP
 from deepmd.jax.common import (
     ArrayAPIVariable,
@@ -12,6 +16,10 @@
 from deepmd.jax.descriptor.base_descriptor import (
     BaseDescriptor,
 )
+from deepmd.jax.env import (
+    flax_version,
+    nnx,
+)
 
 
 @BaseDescriptor.register("hybrid")
@@ -20,7 +28,11 @@ 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

@@ -3,6 +3,10 @@
     Any,
 )
 
+from packaging.version import (
+    Version,
+)
+
 from deepmd.dpmodel.descriptor.repflows import (
     DescrptBlockRepflows as DescrptBlockRepflowsDP,
 )
@@ -12,6 +16,10 @@
     flax_module,
     to_jax_array,
 )
+from deepmd.jax.env import (
+    flax_version,
+    nnx,
+)
 from deepmd.jax.utils.exclude_mask import (
     PairExcludeMask,
 )
@@ -27,8 +35,12 @@ def __setattr__(self, name: str, value: Any) -> None:
             value = to_jax_array(value)
             if value is not None:
                 value = ArrayAPIVariable(value)
+            elif Version(flax_version) >= Version("0.12.0"):
+                value = nnx.data(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 +70,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

@@ -3,6 +3,10 @@
     Any,
 )
 
+from packaging.version import (
+    Version,
+)
+
 from deepmd.dpmodel.descriptor.repformers import (
     Atten2EquiVarApply as Atten2EquiVarApplyDP,
 )
@@ -20,6 +24,10 @@
     flax_module,
     to_jax_array,
 )
+from deepmd.jax.env import (
+    flax_version,
+    nnx,
+)
 from deepmd.jax.utils.exclude_mask import (
     PairExcludeMask,
 )
@@ -36,8 +44,12 @@ def __setattr__(self, name: str, value: Any) -> None:
             value = to_jax_array(value)
             if value is not None:
                 value = ArrayAPIVariable(value)
+            elif Version(flax_version) >= Version("0.12.0"):
+                value = nnx.data(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 +99,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)