fix(pd): add charge_spin stub methods and param to dpa1/dpa2/se_a/se_t_tebd descriptors

Co-authored-by: HydrogenSulfate <23737287+HydrogenSulfate@users.noreply.github.com>

Author: Copilot

deepmd/pd/model/descriptor/dpa1.py

@@ -366,6 +366,18 @@ def get_buffer_type_map(self) -> paddle.Tensor:
         """
         return self.buffer_type_map
 
+    def get_dim_chg_spin(self) -> int:
+        """Returns the dimension of charge_spin input (0 if not supported)."""
+        return 0
+
+    def has_default_chg_spin(self) -> bool:
+        """Returns whether the descriptor has a default charge_spin value."""
+        return False
+
+    def get_default_chg_spin(self) -> None:
+        """Returns the default charge_spin value, or None."""
+        return None
+
     def get_dim_out(self) -> int:
         """Returns the output dimension."""
         ret = self.se_atten.get_dim_out()
@@ -627,6 +639,7 @@ def forward(
         mapping: paddle.Tensor | None = None,
         comm_dict: list[paddle.Tensor] | None = None,
         fparam: paddle.Tensor | None = None,
+        charge_spin: paddle.Tensor | None = None,
     ) -> tuple[
         paddle.Tensor,
         paddle.Tensor | None,

deepmd/pd/model/descriptor/dpa2.py

@@ -333,6 +333,18 @@ def init_subclass_params(sub_data: dict | Any, sub_class: type) -> Any:
             param.stop_gradient = not trainable
         self.compress = False
 
+    def get_dim_chg_spin(self) -> int:
+        """Returns the dimension of charge_spin input (0 if not supported)."""
+        return 0
+
+    def has_default_chg_spin(self) -> bool:
+        """Returns whether the descriptor has a default charge_spin value."""
+        return False
+
+    def get_default_chg_spin(self) -> None:
+        """Returns the default charge_spin value, or None."""
+        return None
+
     def get_rcut(self) -> float:
         """Returns the cut-off radius."""
         return self.rcut
@@ -734,10 +746,8 @@ def forward(
         mapping: paddle.Tensor | None = None,
         comm_dict: list[paddle.Tensor] | None = None,
         fparam: paddle.Tensor | None = None,
+        charge_spin: paddle.Tensor | None = None,
     ) -> tuple[
-        paddle.Tensor,
-        paddle.Tensor | None,
-        paddle.Tensor | None,
         paddle.Tensor | None,
         paddle.Tensor | None,
     ]:

deepmd/pd/model/descriptor/se_a.py

@@ -119,6 +119,18 @@ def __init__(
             seed=seed,
         )
 
+    def get_dim_chg_spin(self) -> int:
+        """Returns the dimension of charge_spin input (0 if not supported)."""
+        return 0
+
+    def has_default_chg_spin(self) -> bool:
+        """Returns whether the descriptor has a default charge_spin value."""
+        return False
+
+    def get_default_chg_spin(self) -> None:
+        """Returns the default charge_spin value, or None."""
+        return None
+
     def get_rcut(self) -> float:
         """Returns the cut-off radius."""
         return self.sea.get_rcut()
@@ -289,6 +301,7 @@ def forward(
         mapping: paddle.Tensor | None = None,
         comm_dict: list[paddle.Tensor] | None = None,
         fparam: paddle.Tensor | None = None,
+        charge_spin: paddle.Tensor | None = None,
     ) -> tuple[
         paddle.Tensor,
         paddle.Tensor | None,

deepmd/pd/model/descriptor/se_t_tebd.py

@@ -190,6 +190,18 @@ def __init__(
         for param in self.parameters():
             param.stop_gradient = not trainable
 
+    def get_dim_chg_spin(self) -> int:
+        """Returns the dimension of charge_spin input (0 if not supported)."""
+        return 0
+
+    def has_default_chg_spin(self) -> bool:
+        """Returns whether the descriptor has a default charge_spin value."""
+        return False
+
+    def get_default_chg_spin(self) -> None:
+        """Returns the default charge_spin value, or None."""
+        return None
+
     def get_rcut(self) -> float:
         """Returns the cut-off radius."""
         return self.se_ttebd.get_rcut()
@@ -438,6 +450,7 @@ def forward(
         mapping: paddle.Tensor | None = None,
         comm_dict: list[paddle.Tensor] | None = None,
         fparam: paddle.Tensor | None = None,
+        charge_spin: paddle.Tensor | None = None,
     ) -> paddle.Tensor:
         """Compute the descriptor.