Merge branch 'devel' into refine_pd_UT

Author: HydrogenSulfate

.github/workflows/test_cuda.yml

@@ -19,7 +19,7 @@ jobs:
     runs-on: nvidia
     # https://github.com/deepmodeling/deepmd-kit/pull/2884#issuecomment-1744216845
     container:
-      image: nvidia/cuda:12.6.2-cudnn-devel-ubuntu22.04
+      image: nvidia/cuda:12.9.1-cudnn-devel-ubuntu22.04
       options: --gpus all
     if: github.repository_owner == 'deepmodeling' && (github.event_name == 'pull_request' && github.event.label && github.event.label.name == 'Test CUDA' || github.event_name == 'workflow_dispatch' || github.event_name == 'merge_group')
     steps:
@@ -49,6 +49,8 @@ jobs:
         export TENSORFLOW_ROOT=$(python -c 'import importlib.util,pathlib;print(pathlib.Path(importlib.util.find_spec("tensorflow").origin).parent)')
         pip install --find-links "https://www.paddlepaddle.org.cn/packages/nightly/cu126/paddlepaddle-gpu/" --index-url https://pypi.org/simple "paddlepaddle-gpu==3.3.0.dev20251204"
         source/install/uv_with_retry.sh pip install --system -v -e .[gpu,test,lmp,cu12,torch,jax] mpi4py --reinstall-package deepmd-kit
+        # See https://github.com/jax-ml/jax/issues/29042
+        source/install/uv_with_retry.sh pip install --system -U 'nvidia-cublas-cu12>=12.9.0.13'
       env:
         DP_VARIANT: cuda
         DP_ENABLE_NATIVE_OPTIMIZATION: 1

deepmd/dpmodel/atomic_model/dp_atomic_model.py

@@ -51,6 +51,8 @@ def __init__(
         self.type_map = type_map
         self.descriptor = descriptor
         self.fitting = fitting
+        if hasattr(self.fitting, "reinit_exclude"):
+            self.fitting.reinit_exclude(self.atom_exclude_types)
         self.type_map = type_map
         super().init_out_stat()
 
@@ -191,7 +193,7 @@ def change_type_map(
             if model_with_new_type_stat is not None
             else None,
         )
-        self.fitting_net.change_type_map(type_map=type_map)
+        self.fitting.change_type_map(type_map=type_map)
 
     def serialize(self) -> dict:
         dd = super().serialize()

deepmd/dpmodel/infer/deep_eval.py

@@ -90,6 +90,7 @@ def __init__(
 
         model_data = load_dp_model(model_file)
         self.dp = BaseModel.deserialize(model_data["model"])
+        self.dp.model_def_script = json.dumps(model_data.get("model_def_script", {}))
         self.rcut = self.dp.get_rcut()
         self.type_map = self.dp.get_type_map()
         if isinstance(auto_batch_size, bool):

deepmd/jax/fitting/fitting.py

@@ -43,6 +43,7 @@ def setattr_for_general_fitting(name: str, value: Any) -> Any:
         "fparam_inv_std",
         "aparam_avg",
         "aparam_inv_std",
+        "case_embd",
         "default_fparam_tensor",
     }:
         value = to_jax_array(value)

deepmd/jax/infer/deep_eval.py

@@ -104,7 +104,7 @@ def __init__(
                 stablehlo_atomic_virial_no_ghost=model_data["@variables"][
                     "stablehlo_atomic_virial_no_ghost"
                 ].tobytes(),
-                model_def_script=model_data["model_def_script"],
+                model_def_script=json.dumps(model_data["model_def_script"]),
                 **model_data["constants"],
             )
         elif model_file.endswith(".savedmodel"):

deepmd/pd/infer/deep_eval.py

@@ -726,7 +726,7 @@ def eval_typeebd(self) -> np.ndarray:
         typeebd = paddle.concat(out, axis=1)
         return to_numpy_array(typeebd)
 
-    def get_model_def_script(self) -> str:
+    def get_model_def_script(self) -> dict:
         """Get model definition script."""
         return self.model_def_script
 

deepmd/pt/infer/deep_eval.py

@@ -684,7 +684,7 @@ def eval_typeebd(self) -> np.ndarray:
         typeebd = torch.cat(out, dim=1)
         return to_numpy_array(typeebd)
 
-    def get_model_def_script(self) -> str:
+    def get_model_def_script(self) -> dict:
         """Get model definition script."""
         return self.model_def_script
 

deepmd/pt/model/atomic_model/dp_atomic_model.py

@@ -64,6 +64,8 @@ def __init__(
         self.rcut = self.descriptor.get_rcut()
         self.sel = self.descriptor.get_sel()
         self.fitting_net = fitting
+        if hasattr(self.fitting_net, "reinit_exclude"):
+            self.fitting_net.reinit_exclude(self.atom_exclude_types)
         super().init_out_stat()
         self.enable_eval_descriptor_hook = False
         self.enable_eval_fitting_last_layer_hook = False
@@ -151,6 +153,9 @@ def change_type_map(
             else None,
         )
         self.fitting_net.change_type_map(type_map=type_map)
+        # Reinitialize fitting to get correct sel_type
+        if hasattr(self.fitting_net, "reinit_exclude"):
+            self.fitting_net.reinit_exclude(self.atom_exclude_types)
 
     def has_message_passing(self) -> bool:
         """Returns whether the atomic model has message passing."""

deepmd/tf/model/model.py

@@ -1003,16 +1003,23 @@ def deserialize(cls, data: dict, suffix: str = "") -> "Descriptor":
         check_version_compatibility(data.pop("@version", 2), 2, 1)
         descriptor = Descriptor.deserialize(data.pop("descriptor"), suffix=suffix)
         # bias_atom_e and out_bias are now completely independent - no conversion needed
-        fitting = Fitting.deserialize(data.pop("fitting"), suffix=suffix)
+        fitting_dict = data.pop("fitting", {})
+        atom_exclude_types = data.pop("atom_exclude_types", [])
+        if len(atom_exclude_types) > 0:
+            # get sel_type from complement of atom_exclude_types
+            full_type_list = np.arange(len(data["type_map"]), dtype=int)
+            sel_type = np.setdiff1d(
+                full_type_list, atom_exclude_types, assume_unique=True
+            )
+            fitting_dict["sel_type"] = sel_type.tolist()
+        fitting = Fitting.deserialize(fitting_dict, suffix=suffix)
         # pass descriptor type embedding to model
         if descriptor.explicit_ntypes:
             type_embedding = descriptor.type_embedding
             fitting.dim_descrpt -= type_embedding.neuron[-1]
         else:
             type_embedding = None
         # BEGINE not supported keys
-        if len(data.pop("atom_exclude_types")) > 0:
-            raise NotImplementedError("atom_exclude_types is not supported")
         if len(data.pop("pair_exclude_types")) > 0:
             raise NotImplementedError("pair_exclude_types is not supported")
         data.pop("rcond", None)

source/tests/array_api_strict/fitting/fitting.py

@@ -31,6 +31,7 @@ def setattr_for_general_fitting(name: str, value: Any) -> Any:
         "fparam_inv_std",
         "aparam_avg",
         "aparam_inv_std",
+        "case_embd",
         "default_fparam_tensor",
     }:
         value = to_array_api_strict_array(value)