Merge branch 'master' into refact-fitting-net

Signed-off-by: Jinzhe Zeng <jinzhe.zeng@ustc.edu.cn>

Author: njzjz

.pre-commit-config.yaml

@@ -63,7 +63,7 @@ repos:
     rev: v21.1.8
     hooks:
       - id: clang-format
-        exclude: ^(source/3rdparty|source/lib/src/gpu/cudart/.+\.inc|.+\.ipynb$|.+\.json$)
+        exclude: ^(source/3rdparty|source/lib/src/gpu/cudart/.+\.inc|.+\.ipynb$|source/tests/infer/.+\.json$)
   # markdown, yaml, CSS, javascript
   - repo: https://github.com/pre-commit/mirrors-prettier
     rev: v4.0.0-alpha.8

deepmd/pt_expt/common.py

@@ -25,10 +25,14 @@
 import numpy as np
 import torch
 
+from deepmd.dpmodel.common import (
+    NativeOP,
+)
+
 # ---------------------------------------------------------------------------
 # dpmodel → pt_expt converter registry
 # ---------------------------------------------------------------------------
-_DPMODEL_TO_PT_EXPT: dict[type, Callable[[Any], torch.nn.Module]] = {}
+_DPMODEL_TO_PT_EXPT: dict[type[NativeOP], Callable[[NativeOP], torch.nn.Module]] = {}
 """Registry mapping dpmodel classes to their pt_expt converter functions.
 
 This registry is populated at module import time via `register_dpmodel_mapping`
@@ -43,7 +47,7 @@
 
 
 def register_dpmodel_mapping(
-    dpmodel_cls: type, converter: Callable[[Any], torch.nn.Module]
+    dpmodel_cls: type[NativeOP], converter: Callable[[NativeOP], torch.nn.Module]
 ) -> None:
     """Register a converter that turns a dpmodel instance into a pt_expt Module.
 
@@ -54,10 +58,10 @@ def register_dpmodel_mapping(
 
     Parameters
     ----------
-    dpmodel_cls : type
+    dpmodel_cls : type[NativeOP]
         The dpmodel class to register (e.g., AtomExcludeMaskDP, NetworkCollectionDP).
         This is the key used for lookup in dpmodel_setattr.
-    converter : Callable[[Any], torch.nn.Module]
+    converter : Callable[[NativeOP], torch.nn.Module]
         A callable that converts a dpmodel instance to a pt_expt module.
         Common patterns:
         - Reconstruct from constructor args: lambda v: PtExptClass(v.ntypes, ...)
@@ -85,7 +89,7 @@ def register_dpmodel_mapping(
 def try_convert_module(value: Any) -> torch.nn.Module | None:
     """Convert a dpmodel object to its pt_expt wrapper if a converter is registered.
 
-    This function looks up the type of *value* in the _DPMODEL_TO_PT_EXPT
+    This function looks up the exact type of *value* in the _DPMODEL_TO_PT_EXPT
     registry. If a converter is found, it invokes it to produce a torch.nn.Module
     wrapper; otherwise it returns None.
 
@@ -103,8 +107,9 @@ def try_convert_module(value: Any) -> torch.nn.Module | None:
 
     Notes
     -----
-    This function uses exact type matching. Each dpmodel class must be explicitly
-    registered via register_dpmodel_mapping.
+    This function uses exact type matching (not isinstance checks) to ensure
+    predictable behavior. Each dpmodel class must be explicitly registered via
+    register_dpmodel_mapping.
 
     The function is called by dpmodel_setattr when it encounters an object that
     might be a dpmodel instance. If conversion succeeds, the caller should use
@@ -211,9 +216,19 @@ def dpmodel_setattr(obj: torch.nn.Module, name: str, value: Any) -> tuple[bool,
 
     # dpmodel object → pt_expt module
     if "_modules" in obj.__dict__:
-        converted = try_convert_module(value)
-        if converted is not None:
-            return False, converted
+        # Try to convert dpmodel objects that aren't already torch.nn.Modules
+        if not isinstance(value, torch.nn.Module):
+            converted = try_convert_module(value)
+            if converted is not None:
+                return False, converted
+            # If this is a NativeOP that should have been registered but wasn't, raise error
+            if isinstance(value, NativeOP):
+                raise TypeError(
+                    f"Attempted to assign a dpmodel object of type {type(value).__name__} "
+                    f"but no converter is registered. Please call register_dpmodel_mapping "
+                    f"for this type. If this object doesn't need conversion, register it "
+                    f"with an identity converter: lambda v: v"
+                )
 
     return False, value
 
@@ -275,3 +290,18 @@ def to_torch_array(array: Any) -> torch.Tensor | None:
     if torch.is_tensor(array):
         return array.to(device=env.DEVICE)
     return torch.as_tensor(array, device=env.DEVICE)
+
+
+# Import utils to trigger dpmodel→pt_expt converter registrations
+# This must happen after the functions above are defined to avoid circular imports
+def _ensure_registrations() -> None:
+    """Import pt_expt.utils modules to register converters.
+
+    This function is called on module import to ensure all dpmodel→pt_expt
+    converters are registered before any descriptors/fittings try to use them.
+    """
+    # Import triggers registration of NetworkCollection, ExcludeMask, EnvMat
+    from deepmd.pt_expt import utils  # noqa: F401
+
+
+_ensure_registrations()

deepmd/pt_expt/utils/__init__.py

@@ -1,5 +1,12 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
 
+from deepmd.dpmodel.utils.env_mat import (
+    EnvMat,
+)
+from deepmd.pt_expt.common import (
+    register_dpmodel_mapping,
+)
+
 from .exclude_mask import (
     AtomExcludeMask,
     PairExcludeMask,
@@ -11,6 +18,10 @@
     TypeEmbedNet,
 )
 
+# Register EnvMat with identity converter - it doesn't need wrapping
+# as it's a stateless utility class
+register_dpmodel_mapping(EnvMat, lambda v: v)
+
 __all__ = [
     "AtomExcludeMask",
     "NetworkCollection",

examples/water_multi_task/pytorch_example/input_torch_with_alias.json

@@ -58,20 +58,31 @@
         "type_map": "type_map_all",
         "descriptor": "dpa3_descriptor",
         "fitting_net": "shared_fit_with_id",
-        "model_branch_alias": ["Default","Water"],
+        "model_branch_alias": [
+          "Default",
+          "Water"
+        ],
         "info": {
-            "description": "Water model with DPA3 descriptor and shared fitting net",
-            "observed_type": ["H", "O"]
+          "description": "Water model with DPA3 descriptor and shared fitting net",
+          "observed_type": [
+            "H",
+            "O"
+          ]
         }
       },
       "water_2": {
         "type_map": "type_map_all",
         "descriptor": "dpa3_descriptor",
         "fitting_net": "shared_fit_with_id",
-        "model_branch_alias": ["Water2"],
+        "model_branch_alias": [
+          "Water2"
+        ],
         "info": {
-            "description": "Water duplicated model with DPA3 descriptor and shared fitting net",
-            "observed_type": ["H", "O"]
+          "description": "Water duplicated model with DPA3 descriptor and shared fitting net",
+          "observed_type": [
+            "H",
+            "O"
+          ]
         }
       }
     }

source/tests/consistent/descriptor/test_dpa1.py

@@ -127,6 +127,7 @@ def data(self) -> dict:
             "type_map": ["O", "H"] if use_econf_tebd else None,
             "seed": 1145141919810,
             "trainable": False,
+            "activation_function": "relu",
         }
 
     def is_meaningless_zero_attention_layer_tests(

source/tests/consistent/descriptor/test_dpa2.py

@@ -132,7 +132,7 @@ def data(self) -> dict:
                     "tebd_dim": 4,
                     "tebd_input_mode": repinit_tebd_input_mode,
                     "set_davg_zero": repinit_set_davg_zero,
-                    "activation_function": "tanh",
+                    "activation_function": "relu",
                     "type_one_side": repinit_type_one_side,
                     "use_three_body": repinit_use_three_body,
                     "three_body_sel": 8,
@@ -163,7 +163,7 @@ def data(self) -> dict:
                     "attn2_hidden": 10,
                     "attn2_nhead": 2,
                     "attn2_has_gate": repformer_attn2_has_gate,
-                    "activation_function": "tanh",
+                    "activation_function": "relu",
                     "update_style": repformer_update_style,
                     "update_residual": 0.001,
                     "update_residual_init": repformer_update_residual_init,

source/tests/consistent/descriptor/test_dpa3.py

@@ -125,7 +125,7 @@ def data(self) -> dict:
                 }
             ),
             # kwargs for descriptor
-            "activation_function": "silu",
+            "activation_function": "relu",
             "precision": precision,
             "exclude_types": exclude_types,
             "env_protection": 0.0,

source/tests/consistent/descriptor/test_hybrid.py

@@ -61,6 +61,7 @@ def data(self) -> dict:
                     "type_one_side": True,
                     "precision": "float64",
                     "seed": 20240229,
+                    "activation_function": "relu",
                 },
                 {
                     "type": "se_e2_a",
@@ -73,6 +74,7 @@ def data(self) -> dict:
                     "type_one_side": True,
                     "precision": "float64",
                     "seed": 20240229,
+                    "activation_function": "relu",
                 },
             ]
         }

source/tests/consistent/descriptor/test_se_atten_v2.py

@@ -124,6 +124,7 @@ def data(self) -> dict:
             "use_tebd_bias": use_tebd_bias,
             "type_map": ["O", "H"] if use_econf_tebd else None,
             "seed": 1145141919810,
+            "activation_function": "relu",
         }
 
     def is_meaningless_zero_attention_layer_tests(

source/tests/consistent/descriptor/test_se_e2_a.py

@@ -99,6 +99,7 @@ def data(self) -> dict:
             "env_protection": env_protection,
             "precision": precision,
             "seed": 1145141919810,
+            "activation_function": "relu",
         }
 
     @property