Merge branch 'deepmodeling:master' into master

Author: zhaiwenxi

deepmd/jax/atomic_model/base_atomic_model.py

@@ -1,35 +1 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
-from typing import (
-    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,
-)
-
-
-def base_atomic_model_set_attr(name: str, value: Any) -> Any:
-    if name in {"out_bias", "out_std"}:
-        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:
-        value = AtomExcludeMask(value.ntypes, value.exclude_types)
-    return value

deepmd/jax/atomic_model/dp_atomic_model.py

@@ -1,12 +1,8 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
-from typing import (
-    Any,
-)
-
+import deepmd.jax.descriptor as _jax_descriptor  # noqa: F401
+import deepmd.jax.fitting.fitting as _jax_fitting  # noqa: F401
+import deepmd.jax.utils.exclude_mask as _jax_exclude_mask  # noqa: F401
 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,
 )
@@ -45,10 +41,6 @@ class jax_atomic_model(dpmodel_atomic_model):
         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)
-
         def forward_common_atomic(
             self,
             extended_coord: jnp.ndarray,

deepmd/jax/atomic_model/linear_atomic_model.py

@@ -3,54 +3,28 @@
     Any,
 )
 
-from packaging.version import (
-    Version,
-)
-
+import deepmd.jax.atomic_model.dp_atomic_model as _jax_dp_atomic_model  # noqa: F401
+import deepmd.jax.atomic_model.pairtab_atomic_model as _jax_pairtab_model  # noqa: F401
+import deepmd.jax.utils.exclude_mask as _jax_exclude_mask  # noqa: F401
 from deepmd.dpmodel.atomic_model.linear_atomic_model import (
     DPZBLLinearEnergyAtomicModel as DPZBLLinearEnergyAtomicModelDP,
 )
-from deepmd.jax.atomic_model.base_atomic_model import (
-    base_atomic_model_set_attr,
-)
-from deepmd.jax.atomic_model.dp_atomic_model import (
-    DPAtomicModel,
-)
-from deepmd.jax.atomic_model.pairtab_atomic_model import (
-    PairTabAtomicModel,
-)
 from deepmd.jax.common import (
-    ArrayAPIVariable,
     flax_module,
-    to_jax_array,
 )
 from deepmd.jax.env import (
-    flax_version,
     jax,
     jnp,
-    nnx,
 )
 
 
 @flax_module
 class DPZBLLinearEnergyAtomicModel(DPZBLLinearEnergyAtomicModelDP):
     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":
+        if name == "zbl_weight":
             # 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

@@ -1,43 +1,19 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
-from typing import (
-    Any,
-)
-
-from packaging.version import (
-    Version,
-)
-
+import deepmd.jax.utils.exclude_mask as _jax_exclude_mask  # noqa: F401
 from deepmd.dpmodel.atomic_model.pairtab_atomic_model import (
     PairTabAtomicModel as PairTabAtomicModelDP,
 )
-from deepmd.jax.atomic_model.base_atomic_model import (
-    base_atomic_model_set_attr,
-)
 from deepmd.jax.common import (
-    ArrayAPIVariable,
     flax_module,
-    to_jax_array,
 )
 from deepmd.jax.env import (
-    flax_version,
     jax,
     jnp,
-    nnx,
 )
 
 
 @flax_module
 class PairTabAtomicModel(PairTabAtomicModelDP):
-    def __setattr__(self, name: str, value: Any) -> None:
-        value = base_atomic_model_set_attr(name, value)
-        if name in {"tab_info", "tab_data"}:
-            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(
         self,
         extended_coord: jnp.ndarray,