fix bug based on coderabbit

Author: ChiahsinChu

deepmd/dpmodel/atomic_model/base_atomic_model.py

@@ -226,6 +226,14 @@ def forward_common_atomic(
             ret_dict[kk] = xp.reshape(tmp_arr, out_shape)
             if atomic_weight is not None:
                 _out_shape = ret_dict[kk].shape
+                # Validate that atomic_weight.shape[2] matches the flattened trailing dimensions
+                # of the output to ensure correct broadcasting behavior
+                out_shape2 = math.prod(_out_shape[2:])
+                if atomic_weight.shape[2] != out_shape2:
+                    raise ValueError(
+                        f"atomic_weight shape {atomic_weight.shape} incompatible with output shape {_out_shape}. "
+                        f"Expected atomic_weight.shape[2] to be {out_shape2} for key '{kk}'"
+                    )
                 ret_dict[kk] = ret_dict[kk] * atomic_weight.reshape(
                     [_out_shape[0], _out_shape[1], -1]
                 )

deepmd/dpmodel/model/make_model.py

@@ -68,7 +68,7 @@ def model_call_from_call_lower(
     fparam: Array | None = None,
     aparam: Array | None = None,
     do_atomic_virial: bool = False,
-    atomic_weight: Optional[np.ndarray] = None,
+    atomic_weight: Array | None = None,
 ) -> dict[str, Array]:
     """Return model prediction from lower interface.
 
@@ -231,7 +231,7 @@ def call(
             fparam: Array | None = None,
             aparam: Array | None = None,
             do_atomic_virial: bool = False,
-            atomic_weight: Optional[np.ndarray] = None,
+            atomic_weight: Array | None = None,
         ) -> dict[str, Array]:
             """Return model prediction.
 
@@ -292,7 +292,7 @@ def call_lower(
             fparam: Array | None = None,
             aparam: Array | None = None,
             do_atomic_virial: bool = False,
-            atomic_weight: Optional[np.ndarray] = None,
+            atomic_weight: Array | None = None,
         ) -> dict[str, Array]:
             """Return model prediction. Lower interface that takes
             extended atomic coordinates and types, nlist, and mapping
@@ -359,7 +359,7 @@ def forward_common_atomic(
             fparam: Array | None = None,
             aparam: Array | None = None,
             do_atomic_virial: bool = False,
-            atomic_weight: Optional[np.ndarray] = None,
+            atomic_weight: Array | None = None,
         ) -> dict[str, Array]:
             atomic_ret = self.atomic_model.forward_common_atomic(
                 extended_coord,
@@ -387,7 +387,9 @@ def input_type_cast(
             fparam: Array | None = None,
             aparam: Array | None = None,
             atomic_weight: Array | None = None,
-        ) -> tuple[Array, Array, np.ndarray | None, np.ndarray | None, np.ndarray | None, str]:
+        ) -> tuple[
+            Array, Array, np.ndarray | None, np.ndarray | None, np.ndarray | None, str
+        ]:
             """Cast the input data to global float type."""
             input_prec = RESERVED_PRECISION_DICT[self.precision_dict[coord.dtype.name]]
             ###

deepmd/pt/infer/deep_eval.py

@@ -379,7 +379,7 @@ def eval(
         else:
             if atomic_weight is not None:
                 raise ValueError(
-                    "atomic_weight is not supported when spin is provided. "
+                    "atomic_weight is not supported when spin is provided."
                 )
             out = self._eval_func(self._eval_model_spin, numb_test, natoms)(
                 coords,
@@ -533,43 +533,9 @@ def _eval_model(
         else:
             aparam_input = None
         if atomic_weight is not None:
-            # Validate atomic_weight shape
-            if atomic_weight.ndim == 1:
-                # (natoms,) - broadcast across frames
-                if atomic_weight.shape[0] != natoms:
-                    raise ValueError(
-                        f"atomic_weight shape {atomic_weight.shape} incompatible with number of atoms {natoms}"
-                    )
-                atomic_weight = np.tile(atomic_weight, (nframes, 1))
-            elif atomic_weight.ndim == 2:
-                # (nframes, natoms) or (natoms, k)
-                if (
-                    atomic_weight.shape[0] == natoms
-                    and atomic_weight.shape[1] != nframes
-                ):
-                    # (natoms, k) - broadcast across frames
-                    atomic_weight = np.tile(atomic_weight, (nframes, 1, 1))
-                elif (
-                    atomic_weight.shape[0] != nframes
-                    or atomic_weight.shape[1] != natoms
-                ):
-                    raise ValueError(
-                        f"atomic_weight shape {atomic_weight.shape} incompatible with nframes={nframes} and natoms={natoms}"
-                    )
-            elif atomic_weight.ndim == 3:
-                # (nframes, natoms, k)
-                if (
-                    atomic_weight.shape[0] != nframes
-                    or atomic_weight.shape[1] != natoms
-                ):
-                    raise ValueError(
-                        f"atomic_weight shape {atomic_weight.shape} incompatible with nframes={nframes} and natoms={natoms}"
-                    )
-            else:
-                raise ValueError(
-                    f"atomic_weight must have 1, 2, or 3 dimensions, got {atomic_weight.ndim}"
-                )
-
+            atomic_weight = self._validate_and_reshape_atomic_weight(
+                atomic_weight, nframes, natoms
+            )
             atomic_weight_input = to_torch_tensor(
                 atomic_weight.reshape(nframes, natoms, -1)
             )
@@ -613,7 +579,6 @@ def _eval_model_spin(
         fparam: np.ndarray | None,
         aparam: np.ndarray | None,
         request_defs: list[OutputVariableDef],
-        atomic_weight: Optional[np.ndarray] = None,
     ) -> tuple[np.ndarray, ...]:
         model = self.dp.to(DEVICE)
 
@@ -655,49 +620,7 @@ def _eval_model_spin(
             )
         else:
             aparam_input = None
-        if atomic_weight is not None:
-            # Validate atomic_weight shape for spin model
-            if atomic_weight.ndim == 1:
-                # (natoms,) - broadcast across frames
-                if atomic_weight.shape[0] != natoms:
-                    raise ValueError(
-                        f"atomic_weight shape {atomic_weight.shape} incompatible with number of atoms {natoms}"
-                    )
-                atomic_weight = np.tile(atomic_weight, (nframes, 1))
-            elif atomic_weight.ndim == 2:
-                # (nframes, natoms) or (natoms, k)
-                if (
-                    atomic_weight.shape[0] == natoms
-                    and atomic_weight.shape[1] != nframes
-                ):
-                    # (natoms, k) - broadcast across frames
-                    atomic_weight = np.tile(atomic_weight, (nframes, 1, 1))
-                elif (
-                    atomic_weight.shape[0] != nframes
-                    or atomic_weight.shape[1] != natoms
-                ):
-                    raise ValueError(
-                        f"atomic_weight shape {atomic_weight.shape} incompatible with nframes={nframes} and natoms={natoms}"
-                    )
-            elif atomic_weight.ndim == 3:
-                # (nframes, natoms, k)
-                if (
-                    atomic_weight.shape[0] != nframes
-                    or atomic_weight.shape[1] != natoms
-                ):
-                    raise ValueError(
-                        f"atomic_weight shape {atomic_weight.shape} incompatible with nframes={nframes} and natoms={natoms}"
-                    )
-            else:
-                raise ValueError(
-                    f"atomic_weight must have 1, 2, or 3 dimensions, got {atomic_weight.ndim}"
-                )
-
-            atomic_weight_input = to_torch_tensor(
-                atomic_weight.reshape(nframes, natoms, -1)
-            )
-        else:
-            atomic_weight_input = None
+        atomic_weight_input = None
 
         do_atomic_virial = any(
             x.category == OutputVariableCategory.DERV_C_REDU for x in request_defs
@@ -954,3 +877,60 @@ def eval_fitting_last_layer(
         fitting_net = model.eval_fitting_last_layer()
         model.set_eval_fitting_last_layer_hook(False)
         return to_numpy_array(fitting_net)
+
+    def _validate_and_reshape_atomic_weight(
+        self,
+        atomic_weight: np.ndarray,
+        nframes: int,
+        natoms: int,
+    ) -> np.ndarray:
+        """Validate and reshape atomic_weight to (nframes, natoms, k).
+
+        Parameters
+        ----------
+        atomic_weight : np.ndarray
+            Input atomic weights with various possible shapes
+        nframes : int
+            Number of frames
+        natoms : int
+            Number of atoms
+
+        Returns
+        -------
+        np.ndarray
+            Reshaped atomic_weight with shape (nframes, natoms, k)
+
+        Raises
+        ------
+        ValueError
+            If atomic_weight shape is incompatible
+        """
+        # Validate atomic_weight shape
+        if atomic_weight.ndim == 1:
+            # (natoms,) - broadcast across frames
+            if atomic_weight.shape[0] != natoms:
+                raise ValueError(
+                    f"atomic_weight shape {atomic_weight.shape} incompatible with number of atoms {natoms}"
+                )
+            atomic_weight = np.tile(atomic_weight, (nframes, 1))
+        elif atomic_weight.ndim == 2:
+            # (nframes, natoms) or (natoms, k)
+            if atomic_weight.shape[0] == natoms and atomic_weight.shape[1] != nframes:
+                # (natoms, k) - broadcast across frames
+                atomic_weight = np.tile(atomic_weight, (nframes, 1, 1))
+            elif atomic_weight.shape[0] != nframes or atomic_weight.shape[1] != natoms:
+                raise ValueError(
+                    f"atomic_weight shape {atomic_weight.shape} incompatible with nframes={nframes} and natoms={natoms}"
+                )
+        elif atomic_weight.ndim == 3:
+            # (nframes, natoms, k)
+            if atomic_weight.shape[0] != nframes or atomic_weight.shape[1] != natoms:
+                raise ValueError(
+                    f"atomic_weight shape {atomic_weight.shape} incompatible with nframes={nframes} and natoms={natoms}"
+                )
+        else:
+            raise ValueError(
+                f"atomic_weight must have 1, 2, or 3 dimensions, got {atomic_weight.ndim}"
+            )
+
+        return atomic_weight

deepmd/pt/model/model/dipole_model.py

@@ -60,7 +60,7 @@ def forward(
         fparam: torch.Tensor | None = None,
         aparam: torch.Tensor | None = None,
         do_atomic_virial: bool = False,
-        atomic_weight: Optional[torch.Tensor] = None,
+        atomic_weight: torch.Tensor | None = None,
     ) -> dict[str, torch.Tensor]:
         model_ret = self.forward_common(
             coord,

deepmd/pt/model/model/make_model.py

@@ -138,7 +138,7 @@ def forward_common(
             fparam: torch.Tensor | None = None,
             aparam: torch.Tensor | None = None,
             do_atomic_virial: bool = False,
-            atomic_weight: Optional[torch.Tensor] = None,
+            atomic_weight: torch.Tensor | None = None,
         ) -> dict[str, torch.Tensor]:
             """Return model prediction.
 
@@ -255,7 +255,7 @@ def forward_common_lower(
             do_atomic_virial: bool = False,
             comm_dict: dict[str, torch.Tensor] | None = None,
             extra_nlist_sort: bool = False,
-            atomic_weight: Optional[torch.Tensor] = None,
+            atomic_weight: torch.Tensor | None = None,
         ) -> dict[str, torch.Tensor]:
             """Return model prediction. Lower interface that takes
             extended atomic coordinates and types, nlist, and mapping
@@ -642,7 +642,7 @@ def forward(
             fparam: torch.Tensor | None = None,
             aparam: torch.Tensor | None = None,
             do_atomic_virial: bool = False,
-            atomic_weight: Optional[torch.Tensor] = None,
+            atomic_weight: torch.Tensor | None = None,
         ) -> dict[str, torch.Tensor]:
             # directly call the forward_common method when no specific transform rule
             return self.forward_common(