Skip to content

fix(tf): handle tensor model bias reshape in backend conversion #4957

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

Closed
Copilot wants to merge 5 commits into from
Closed

fix(tf): handle tensor model bias reshape in backend conversion #4957

Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
ac6b184
Initial plan
Copilot Sep 3, 2025
3a79386
fix(tf): handle dipole model bias reshape in backend conversion
Copilot Sep 3, 2025
eaf5411
fix(tf): extend bias reshape fix to cover all fitting types
Copilot Sep 4, 2025
95aef5d
Update model.py
njzjz Sep 4, 2025
afc4c38
fix serialize
njzjz Sep 5, 2025
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
42 changes: 29 additions & 13 deletions deepmd/tf/model/model.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -826,19 +826,32 @@ def deserialize(cls, data: dict, suffix: str = "") -> "Descriptor":
# deepcopy is not used for performance reasons
data["fitting"] = data["fitting"].copy()
data["fitting"]["@variables"] = data["fitting"]["@variables"].copy()
if (
int(np.any(data["fitting"]["@variables"]["bias_atom_e"]))
+ int(np.any(data["@variables"]["out_bias"]))
> 1
):
raise ValueError(
"fitting/@variables/bias_atom_e and @variables/out_bias should not be both non-zero"

# For InvarFitting types (ener, dos, property), out_bias can be reshaped and added to bias_atom_e
# For GeneralFitting types (dipole, polar), out_bias and bias_atom_e have different purposes and shapes
# and should not be added together
fitting_type = data["fitting"].get("type", "ener")
if fitting_type in [
"ener",
"dos",
"property",
]:
# For InvarFitting types, use the original logic to reshape and add out_bias to bias_atom_e
if (
int(np.any(data["fitting"]["@variables"]["bias_atom_e"]))
+ int(np.any(data["@variables"]["out_bias"]))
> 1
):
raise ValueError(
"fitting/@variables/bias_atom_e and @variables/out_bias should not be both non-zero"
)
data["fitting"]["@variables"]["bias_atom_e"] = data["fitting"][
"@variables"
]["bias_atom_e"] + data["@variables"]["out_bias"].reshape(
data["fitting"]["@variables"]["bias_atom_e"].shape
)
Comment thread
njzjz marked this conversation as resolved.
data["fitting"]["@variables"]["bias_atom_e"] = data["fitting"][
"@variables"
]["bias_atom_e"] + data["@variables"]["out_bias"].reshape(
data["fitting"]["@variables"]["bias_atom_e"].shape
)
# For GeneralFitting types (dipole, polar), keep out_bias separate - don't add to bias_atom_e
# These fitting types have different bias structures that are incompatible
fitting = Fitting.deserialize(data.pop("fitting"), suffix=suffix)
# pass descriptor type embedding to model
if descriptor.explicit_ntypes:
Expand Down Expand Up @@ -887,7 +900,10 @@ def serialize(self, suffix: str = "") -> dict:

ntypes = len(self.get_type_map())
dict_fit = self.fitting.serialize(suffix=suffix)
if dict_fit.get("@variables", {}).get("bias_atom_e") is not None:
if (
dict_fit.get("@variables", {}).get("bias_atom_e") is not None
and dict_fit["dim_out"] == dict_fit["embedding_width"]
):
out_bias = dict_fit["@variables"]["bias_atom_e"].reshape(
[1, ntypes, dict_fit["dim_out"]]
)
Expand Down
Loading