Skip to content

fix(pd): add default chg spin parameter #5469

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
HydrogenSulfate wants to merge 4 commits into
base: master
Choose a base branch
from
+99 −40
Open

fix(pd): add default chg spin parameter #5469

Show file tree
Hide file tree
Changes from 3 commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
fe9f626
Sync Paddle DPA3 with dpmodel: add default_chg_spin, rename fparam to…
Copilot May 27, 2026
0e792ac
fix(pd): add charge_spin stub methods and param to dpa1/dpa2/se_a/se_…
Copilot May 27, 2026
db26f5e
fix(pd): update dp_atomic_model to call descriptor forward with charg…
Copilot May 27, 2026
6ef79f2
test(pd): sync unit tests with charge_spin interface changes
Copilot Jun 1, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion deepmd/pd/model/atomic_model/dp_atomic_model.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -332,7 +332,7 @@ def forward_atomic(
nlist,
mapping=mapping,
comm_dict=comm_dict,
fparam=fparam if self.add_chg_spin_ebd else None,
charge_spin=fparam if self.add_chg_spin_ebd else None,
Copy link
Copy Markdown
Contributor

@coderabbitai coderabbitai Bot May 27, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟠 Major | ⚡ Quick win

Missing default charge_spin fallback breaks the new default contract.

When self.add_chg_spin_ebd is enabled and fparam is None, this passes charge_spin=None, which trips the DPA3 assertion path and ignores configured default_chg_spin.

💡 Suggested fix 
         nframes, nloc, nnei = nlist.shape
         atype = extended_atype[:, :nloc]
         if self.do_grad_r() or self.do_grad_c():
             extended_coord.stop_gradient = False
+        charge_spin = fparam if self.add_chg_spin_ebd else None
+        if self.add_chg_spin_ebd and charge_spin is None:
+            default_cs = self.descriptor.get_default_chg_spin()
+            if default_cs is not None:
+                cs = paddle.to_tensor(default_cs, dtype=extended_coord.dtype).to(
+                    device=extended_coord.place
+                )
+                charge_spin = paddle.tile(cs.reshape([1, -1]), [nframes, 1])
         descriptor, rot_mat, g2, h2, sw = self.descriptor(
             extended_coord,
             extended_atype,
             nlist,
             mapping=mapping,
             comm_dict=comm_dict,
-            charge_spin=fparam if self.add_chg_spin_ebd else None,
+            charge_spin=charge_spin,
         )
🤖 Prompt for AI Agents 
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@deepmd/pd/model/atomic_model/dp_atomic_model.py` at line 335, When
constructing the call that sets charge_spin (the line currently using
"charge_spin=fparam if self.add_chg_spin_ebd else None"), ensure that when
self.add_chg_spin_ebd is True but fparam is None you fall back to the configured
default_chg_spin instead of passing None; change the expression to use fparam if
fparam is not None else self.default_chg_spin (only when self.add_chg_spin_ebd
is True), so charge_spin becomes (fparam if fparam is not None else
self.default_chg_spin) when self.add_chg_spin_ebd else None and avoids
triggering the DPA3 assertion path.

)
assert descriptor is not None
if self.enable_eval_descriptor_hook:
Expand Down
13 changes: 13 additions & 0 deletions deepmd/pd/model/descriptor/dpa1.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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()
Expand Down Expand Up @@ -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,
Expand Down
16 changes: 13 additions & 3 deletions deepmd/pd/model/descriptor/dpa2.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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,
]:
Expand Down
27 changes: 23 additions & 4 deletions deepmd/pd/model/descriptor/dpa3.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -122,6 +122,7 @@ def __init__(
use_loc_mapping: bool = True,
type_map: list[str] | None = None,
add_chg_spin_ebd: bool = False,
default_chg_spin: list[float] | None = None,
) -> None:
super().__init__()

Expand Down Expand Up @@ -177,6 +178,11 @@ def init_subclass_params(sub_data: dict | Any, sub_class: type) -> Any:

self.use_econf_tebd = use_econf_tebd
self.add_chg_spin_ebd = add_chg_spin_ebd
if default_chg_spin is not None and len(default_chg_spin) != 2:
raise ValueError(
"default_chg_spin must have exactly 2 values [charge, spin]"
)
self.default_chg_spin = default_chg_spin
self.use_loc_mapping = use_loc_mapping
self.use_tebd_bias = use_tebd_bias
self.type_map = type_map
Expand Down Expand Up @@ -447,6 +453,18 @@ def get_stat_mean_and_stddev(
stddev_list = [self.repflows.stddev]
return mean_list, stddev_list

def get_dim_chg_spin(self) -> int:
"""Returns the dimension of charge_spin input."""
return 2 if self.add_chg_spin_ebd else 0

def has_default_chg_spin(self) -> bool:
"""Returns whether default charge_spin values are set."""
return self.default_chg_spin is not None

def get_default_chg_spin(self) -> list[float] | None:
"""Returns the default charge_spin values."""
return self.default_chg_spin

def serialize(self) -> dict:
repflows = self.repflows
data = {
Expand All @@ -465,6 +483,7 @@ def serialize(self) -> dict:
"use_tebd_bias": self.use_tebd_bias,
"use_loc_mapping": self.use_loc_mapping,
"add_chg_spin_ebd": self.add_chg_spin_ebd,
"default_chg_spin": self.default_chg_spin,
"type_map": self.type_map,
"type_embedding": self.type_embedding.embedding.serialize(),
}
Expand Down Expand Up @@ -541,7 +560,7 @@ def forward(
nlist: paddle.Tensor,
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,
Expand Down Expand Up @@ -593,11 +612,11 @@ def forward(
node_ebd_ext = self.type_embedding(extended_atype)

if self.add_chg_spin_ebd:
assert fparam is not None
assert charge_spin is not None
assert self.chg_embedding is not None
assert self.spin_embedding is not None
charge = fparam[:, 0].to(dtype=paddle.int64) + 100
spin = fparam[:, 1].to(dtype=paddle.int64)
charge = charge_spin[:, 0].to(dtype=paddle.int64) + 100
spin = charge_spin[:, 1].to(dtype=paddle.int64)
chg_ebd = self.chg_embedding(charge)
spin_ebd = self.spin_embedding(spin)
sys_cs_embd = self.act(
Expand Down
13 changes: 13 additions & 0 deletions deepmd/pd/model/descriptor/se_a.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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()
Expand Down Expand Up @@ -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,
Expand Down
13 changes: 13 additions & 0 deletions deepmd/pd/model/descriptor/se_t_tebd.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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()
Expand Down Expand Up @@ -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.

Expand Down
Loading