From 6194ff62cf0ae2ac236b0936de4287df8daae76f Mon Sep 17 00:00:00 2001 From: Jinzhe Zeng Date: Wed, 8 Jul 2026 00:48:48 +0800 Subject: [PATCH 1/4] feat(tf2): support DPA4 training --- deepmd/dpmodel/array_api.py | 24 ++ deepmd/tf2/common.py | 142 ++++++++++ deepmd/tf2/descriptor/__init__.py | 4 + deepmd/tf2/descriptor/dpa4.py | 267 ++++++++++++++++++ deepmd/tf2/fitting/__init__.py | 4 + deepmd/tf2/fitting/dpa4_ener.py | 53 ++++ deepmd/tf2/model/ener_model.py | 2 + deepmd/tf2/model/model.py | 49 ++++ deepmd/tf2/train/trainer.py | 50 +++- .../tests/consistent/descriptor/test_dpa4.py | 17 ++ .../consistent/fitting/test_dpa4_ener.py | 21 ++ 11 files changed, 631 insertions(+), 2 deletions(-) create mode 100644 deepmd/tf2/descriptor/dpa4.py create mode 100644 deepmd/tf2/fitting/dpa4_ener.py diff --git a/deepmd/dpmodel/array_api.py b/deepmd/dpmodel/array_api.py index 115242edfb..4e2ace9bbe 100644 --- a/deepmd/dpmodel/array_api.py +++ b/deepmd/dpmodel/array_api.py @@ -201,6 +201,18 @@ def xp_add_at(x: Array, indices: Array, values: Array) -> Array: import torch return torch.index_add(x, 0, indices, values) + elif getattr(xp, "__name__", "") == "deepmd._vendors.ndtensorflow": + import tensorflow as tf + + x_tensor = x.unwrap() + indices_tensor = tf.reshape(tf.cast(indices.unwrap(), tf.int64), (-1, 1)) + values_tensor = values.unwrap() + updates = tf.scatter_nd( + indices_tensor, + values_tensor, + tf.shape(x_tensor, out_type=tf.int64), + ) + return xp.asarray(x_tensor + updates) else: # Fallback for array_api_strict: use basic indexing only # may need a more efficient way to do this @@ -270,6 +282,18 @@ def xp_maximum_at(x: Array, indices: Array, values: Array) -> Array: return torch.scatter_reduce( x, 0, index, values, reduce="amax", include_self=True ) + elif getattr(xp, "__name__", "") == "deepmd._vendors.ndtensorflow": + import tensorflow as tf + + x_tensor = x.unwrap() + indices_tensor = tf.reshape(tf.cast(indices.unwrap(), tf.int64), (-1,)) + values_tensor = values.unwrap() + reduced = tf.math.unsorted_segment_max( + values_tensor, + indices_tensor, + tf.shape(x_tensor, out_type=tf.int64)[0], + ) + return xp.asarray(tf.maximum(x_tensor, reduced)) else: # Fallback for array_api_strict: basic indexing only. n = indices.shape[0] diff --git a/deepmd/tf2/common.py b/deepmd/tf2/common.py index bb8155a38c..f3dfc8b12a 100644 --- a/deepmd/tf2/common.py +++ b/deepmd/tf2/common.py @@ -2,6 +2,8 @@ from collections.abc import ( Callable, + Mapping, + Sequence, ) from functools import ( wraps, @@ -104,6 +106,7 @@ def unwrap_value(value: Any) -> Any: f"{_PACKAGE_ROOT}.descriptor.dpa2", f"{_PACKAGE_ROOT}.descriptor.repflows", f"{_PACKAGE_ROOT}.descriptor.dpa3", + f"{_PACKAGE_ROOT}.descriptor.dpa4", f"{_PACKAGE_ROOT}.descriptor.hybrid", f"{_PACKAGE_ROOT}.fitting", f"{_PACKAGE_ROOT}.atomic_model.dp_atomic_model", @@ -253,6 +256,47 @@ def tf2_module(module: type[T]) -> type[T]: @wraps(module, updated=()) class TF2Module(module, tf.Module): # type: ignore[misc, valid-type] + @staticmethod + def _tf2_array_variable_storage_name(name: str) -> str: + return f"_tf2_{name}_variable" + + @staticmethod + def _tf2_array_variable_list_storage_name(name: str) -> str: + return f"_tf2_{name}_variables" + + def _tf2_array_variable_attr_names(self) -> set[str]: + return set(getattr(self, "_tf2_array_variable_attrs", ())) + + def _tf2_array_variable_list_attr_names(self) -> set[str]: + return set(getattr(self, "_tf2_array_variable_list_attrs", ())) + + def _set_tf2_array_variable(self, name: str, value: Any) -> None: + storage_name = self._tf2_array_variable_storage_name(name) + if value is None: + tf.Module.__setattr__(self, storage_name, None) + return + tensor = to_tf_tensor(value) + variable = tf.Variable( + tensor, + trainable=bool(getattr(self, "trainable", True)), + name=name, + ) + tf.Module.__setattr__(self, storage_name, variable) + + def _set_tf2_array_variable_list(self, name: str, value: Any) -> None: + storage_name = self._tf2_array_variable_list_storage_name(name) + variables = [] + for idx, item in enumerate(value): + tensor = to_tf_tensor(item) + variables.append( + tf.Variable( + tensor, + trainable=bool(getattr(self, "trainable", True)), + name=f"{name}_{idx}", + ) + ) + tf.Module.__setattr__(self, storage_name, variables) + def __init__(self, *args: Any, **kwargs: Any) -> None: tf.Module.__init__(self) super().__init__(*args, **kwargs) @@ -266,11 +310,109 @@ def __init__(self, *args: Any, **kwargs: Any) -> None: ) if converted is not value: setattr(self, name, converted) + self._refresh_tf2_trackable_lists() + + def _refresh_tf2_trackable_lists(self) -> None: + """Rebuild trackable list containers after backend conversion.""" + seen: set[int] = set() + + def visit(value: Any) -> None: + if value is None or isinstance(value, (str, bytes, int, float, bool)): + return + if isinstance(value, (np.ndarray, tf.Tensor, tf.Variable, xp.Array)): + return + value_id = id(value) + if value_id in seen: + return + seen.add(value_id) + + if isinstance(value, Mapping): + for item in value.values(): + visit(item) + return + if isinstance(value, Sequence) and not isinstance(value, (str, bytes)): + for item in value: + visit(item) + return + + try: + value_dict = object.__getattribute__(value, "__dict__") + except AttributeError: + return + + for attr_name, attr_value in list(value_dict.items()): + if attr_name.startswith("_"): + continue + if not isinstance(attr_value, list): + continue + if any(isinstance(item, tf.Module) for item in attr_value): + setattr(value, attr_name, list(attr_value)) + + try: + value_dict = object.__getattribute__(value, "__dict__") + except AttributeError: + return + for attr_name, attr_value in list(value_dict.items()): + if attr_name.startswith("_"): + continue + visit(attr_value) + + visit(self) + + def __getattribute__(self, name: str) -> Any: + if not name.startswith("_tf2_"): + array_attrs = object.__getattribute__( + self, + "_tf2_array_variable_attr_names", + )() + if name in array_attrs: + storage_name = object.__getattribute__( + self, + "_tf2_array_variable_storage_name", + )(name) + variable = object.__getattribute__(self, storage_name) + return None if variable is None else to_tensorflow_array(variable) + + list_attrs = object.__getattribute__( + self, + "_tf2_array_variable_list_attr_names", + )() + if name in list_attrs: + storage_name = object.__getattribute__( + self, + "_tf2_array_variable_list_storage_name", + )(name) + variables = object.__getattribute__(self, storage_name) + return [to_tensorflow_array(var) for var in variables] + return super().__getattribute__(name) def __setattr__(self, name: str, value: Any) -> None: + if name in self._tf2_array_variable_attr_names(): + self._set_tf2_array_variable(name, value) + return + if name in self._tf2_array_variable_list_attr_names(): + self._set_tf2_array_variable_list(name, value) + return value = tf2_setattr(self, name, value) return super().__setattr__(name, value) + original_deserialize = getattr(module, "deserialize", None) + if original_deserialize is not None: + + @classmethod + def deserialize(cls: type[Any], data: Any) -> Any: + deserialize_func = getattr(original_deserialize, "__func__", None) + if deserialize_func is None: + obj = original_deserialize(data) + else: + obj = deserialize_func(cls, data) + refresh = getattr(obj, "_refresh_tf2_trackable_lists", None) + if callable(refresh): + refresh() + return obj + + TF2Module.deserialize = deserialize + if hasattr(TF2Module, "deserialize"): for base in module.__bases__: if base in (object, NativeOP): diff --git a/deepmd/tf2/descriptor/__init__.py b/deepmd/tf2/descriptor/__init__.py index 1bbefbea6f..b9235aa8b5 100644 --- a/deepmd/tf2/descriptor/__init__.py +++ b/deepmd/tf2/descriptor/__init__.py @@ -8,6 +8,9 @@ from .dpa3 import ( DescrptDPA3, ) +from .dpa4 import ( + DescrptDPA4, +) from .hybrid import ( DescrptHybrid, ) @@ -31,6 +34,7 @@ "DescrptDPA1", "DescrptDPA2", "DescrptDPA3", + "DescrptDPA4", "DescrptHybrid", "DescrptSeA", "DescrptSeAttenV2", diff --git a/deepmd/tf2/descriptor/dpa4.py b/deepmd/tf2/descriptor/dpa4.py new file mode 100644 index 0000000000..b706cac4f9 --- /dev/null +++ b/deepmd/tf2/descriptor/dpa4.py @@ -0,0 +1,267 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +from collections.abc import ( + Mapping, + Sequence, +) +from typing import ( + Any, +) + +import numpy as np + +from deepmd._vendors import ndtensorflow as xp +from deepmd.dpmodel.common import ( + NativeOP, +) +from deepmd.dpmodel.descriptor.dpa4 import DescrptDPA4 as DescrptDPA4DP +from deepmd.dpmodel.descriptor.dpa4_nn.activation import SwiGLU as SwiGLUDP +from deepmd.dpmodel.descriptor.dpa4_nn.grid_net import GridProduct as GridProductDP +from deepmd.dpmodel.descriptor.dpa4_nn.radial import ( + C3CutoffEnvelope as C3CutoffEnvelopeDP, +) +from deepmd.dpmodel.descriptor.dpa4_nn.radial import RadialMLP as RadialMLPDP +from deepmd.dpmodel.descriptor.dpa4_nn.so2 import SO2Linear as SO2LinearDP +from deepmd.dpmodel.descriptor.dpa4_nn.wignerd import ( + WignerDCalculator as WignerDCalculatorDP, +) +from deepmd.tf2.common import ( + register_dpmodel_mapping, + tf, + tf2_module, + to_tf_tensor, + try_convert_module, +) +from deepmd.tf2.descriptor.base_descriptor import ( + BaseDescriptor, +) +from deepmd.tf2.utils import exclude_mask as _tf2_exclude_mask # noqa: F401 +from deepmd.tf2.utils import network as _tf2_network # noqa: F401 + + +@tf2_module +class SwiGLU(SwiGLUDP): + pass + + +register_dpmodel_mapping(SwiGLUDP, lambda v: SwiGLU()) + + +@tf2_module +class C3CutoffEnvelope(C3CutoffEnvelopeDP): + pass + + +register_dpmodel_mapping( + C3CutoffEnvelopeDP, + lambda v: C3CutoffEnvelope(v.rcut, v.p, precision=v.precision), +) + + +@tf2_module +class RadialMLP(RadialMLPDP): + def __init__(self, *args: Any, **kwargs: Any) -> None: + super().__init__(*args, **kwargs) + self.net = [self._convert_layer(layer) for layer in self.net] + self._tracked_net_modules = [ + layer for layer in self.net if isinstance(layer, tf.Module) + ] + + @staticmethod + def _convert_layer(layer: Any) -> Any: + if isinstance(layer, tf.Module): + return layer + if isinstance(layer, NativeOP): + converted = try_convert_module(layer) + if converted is not None: + return converted + return layer + + +register_dpmodel_mapping( + RadialMLPDP, + lambda v: RadialMLP.deserialize(v.serialize()), +) + + +@tf2_module +class GridProduct(GridProductDP): + pass + + +register_dpmodel_mapping(GridProductDP, lambda v: GridProduct()) + + +@tf2_module +class WignerDCalculator(WignerDCalculatorDP): + pass + + +register_dpmodel_mapping( + WignerDCalculatorDP, + lambda v: WignerDCalculator(v.lmax, eps=v.eps, precision=v.precision), +) + + +_TRAINABLE_ATTRS: dict[str, tuple[str, ...]] = { + "RMSNorm": ("adam_scale",), + "EquivariantRMSNorm": ("adam_scale", "bias"), + "ReducedEquivariantRMSNorm": ("adam_scale", "bias0"), + "ScalarRMSNorm": ("adam_scale",), + "RadialBasis": ("adam_freqs",), + "SO3Linear": ("weight", "bias"), + "FocusLinear": ("weight", "bias"), + "ChannelLinear": ("weight", "bias"), + "SO2Linear": ("weight_m0", "bias0"), + "DynamicRadialDegreeMixer": ("weight", "channel_basis"), + "SO2Convolution": ( + "adamw_attn_logit_w", + "adamw_attn_z_bias_raw", + "adamw_attn_gate_w", + "adamw_focus_compete_w", + "focus_compete_bias", + ), + "SeZMTypeEmbedding": ("adam_type_embedding",), + "SpinEmbedding": ("adam_spin_vec_weight", "adam_spin_nbr_weight"), + "EnvironmentInitialEmbedding": ("spin_scale",), + "DepthAttnRes": ("adamw_pseudo_query",), + "S2GridNet": ("residual_scale",), + "SO3GridNet": ("residual_scale",), + "DescrptDPA4": ("film_scale_strength_log", "film_shift_strength_log"), +} + +_TRAINABLE_LIST_ATTRS: dict[str, tuple[str, ...]] = { + "SO2Linear": ("weight_m",), + "SO2Convolution": ("adam_so2_layer_scales",), +} + + +def _is_array_like(value: Any) -> bool: + return isinstance(value, (np.ndarray, tf.Tensor, tf.Variable, xp.Array)) + + +def _is_floating_array(value: Any) -> bool: + tensor = to_tf_tensor(value) + return tensor is not None and tensor.dtype.is_floating + + +def _iter_object_tree(root: Any) -> Any: + seen: set[int] = set() + + def visit(value: Any) -> Any: + if value is None or isinstance(value, (str, bytes, int, float, bool)): + return + if _is_array_like(value): + return + value_id = id(value) + if value_id in seen: + return + seen.add(value_id) + + if isinstance(value, Mapping): + for item in value.values(): + yield from visit(item) + return + if isinstance(value, Sequence): + for item in value: + yield from visit(item) + return + try: + value_dict = object.__getattribute__(value, "__dict__") + except AttributeError: + return + + yield value + for item in value_dict.values(): + yield from visit(item) + + yield from visit(root) + + +def _enable_tf2_array_variable_attr(module: Any, name: str) -> None: + attrs = set(getattr(module, "_tf2_array_variable_attrs", ())) + if name not in attrs: + tf.Module.__setattr__(module, "_tf2_array_variable_attrs", attrs | {name}) + + +def _enable_tf2_array_variable_list_attr(module: Any, name: str) -> None: + attrs = set(getattr(module, "_tf2_array_variable_list_attrs", ())) + if name not in attrs: + tf.Module.__setattr__( + module, + "_tf2_array_variable_list_attrs", + attrs | {name}, + ) + + +def _promote_trainable(module: Any, names: tuple[str, ...]) -> None: + if not getattr(module, "trainable", True): + return + for name in names: + if not hasattr(module, name): + continue + value = getattr(module, name) + if not _is_floating_array(value): + continue + _enable_tf2_array_variable_attr(module, name) + setattr(module, name, value) + + +def _promote_trainable_lists(module: Any, names: tuple[str, ...]) -> None: + if not getattr(module, "trainable", True): + return + for name in names: + if not hasattr(module, name): + continue + value = getattr(module, name) + if not isinstance(value, Sequence) or isinstance(value, (str, bytes)): + continue + if not value or not all(_is_floating_array(item) for item in value): + continue + _enable_tf2_array_variable_list_attr(module, name) + setattr(module, name, value) + + +def _promote_trainable_tree(module: Any) -> Any: + for submodule in _iter_object_tree(module): + names = _TRAINABLE_ATTRS.get(type(submodule).__name__) + if names is not None: + _promote_trainable(submodule, names) + list_names = _TRAINABLE_LIST_ATTRS.get(type(submodule).__name__) + if list_names is not None: + _promote_trainable_lists(submodule, list_names) + return module + + +@tf2_module +class SO2Linear(SO2LinearDP): + def __init__(self, *args: Any, **kwargs: Any) -> None: + super().__init__(*args, **kwargs) + _promote_trainable_lists(self, ("weight_m",)) + + @classmethod + def deserialize(cls, data: dict) -> "SO2Linear": + obj = super().deserialize(data) + _promote_trainable_lists(obj, ("weight_m",)) + return obj + + +register_dpmodel_mapping( + SO2LinearDP, + lambda v: SO2Linear.deserialize(v.serialize()), +) + + +@BaseDescriptor.register("SeZM") +@BaseDescriptor.register("sezm") +@BaseDescriptor.register("DPA4") +@BaseDescriptor.register("dpa4") +@tf2_module +class DescrptDPA4(DescrptDPA4DP): + def __init__(self, *args: Any, **kwargs: Any) -> None: + super().__init__(*args, **kwargs) + _promote_trainable_tree(self) + + @classmethod + def deserialize(cls, data: dict) -> "DescrptDPA4": + obj = super().deserialize(data) + return _promote_trainable_tree(obj) diff --git a/deepmd/tf2/fitting/__init__.py b/deepmd/tf2/fitting/__init__.py index 2041f600ea..81aee58175 100644 --- a/deepmd/tf2/fitting/__init__.py +++ b/deepmd/tf2/fitting/__init__.py @@ -1,4 +1,7 @@ # SPDX-License-Identifier: LGPL-3.0-or-later +from .dpa4_ener import ( + SeZMEnergyFittingNet, +) from .fitting import ( DipoleFittingNet, DOSFittingNet, @@ -13,4 +16,5 @@ "EnergyFittingNet", "PolarFittingNet", "PropertyFittingNet", + "SeZMEnergyFittingNet", ] diff --git a/deepmd/tf2/fitting/dpa4_ener.py b/deepmd/tf2/fitting/dpa4_ener.py new file mode 100644 index 0000000000..d20e06060c --- /dev/null +++ b/deepmd/tf2/fitting/dpa4_ener.py @@ -0,0 +1,53 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +from typing import ( + Any, + ClassVar, +) + +from deepmd.dpmodel.fitting.dpa4_ener import GLUFittingNet as GLUFittingNetDP +from deepmd.dpmodel.fitting.dpa4_ener import ( + SeZMEnergyFittingNet as SeZMEnergyFittingNetDP, +) +from deepmd.dpmodel.fitting.dpa4_ener import ( + SeZMNetworkCollection as SeZMNetworkCollectionDP, +) +from deepmd.tf2.common import ( + register_dpmodel_mapping, + tf2_module, +) +from deepmd.tf2.fitting.base_fitting import ( + BaseFitting, +) +from deepmd.tf2.utils import network as _tf2_network # noqa: F401 + + +@tf2_module +class GLUFittingNet(GLUFittingNetDP): + pass + + +register_dpmodel_mapping( + GLUFittingNetDP, + lambda v: GLUFittingNet.deserialize(v.serialize()), +) + + +@tf2_module +class SeZMNetworkCollection(SeZMNetworkCollectionDP): + NETWORK_TYPE_MAP: ClassVar[dict[str, type]] = { + "sezm_fitting_network": GLUFittingNet, + } + + +register_dpmodel_mapping( + SeZMNetworkCollectionDP, + lambda v: SeZMNetworkCollection.deserialize(v.serialize()), +) + + +@BaseFitting.register("dpa4_ener") +@BaseFitting.register("sezm_ener") +@tf2_module +class SeZMEnergyFittingNet(SeZMEnergyFittingNetDP): + def __setattr__(self, name: str, value: Any) -> None: + return super().__setattr__(name, value) diff --git a/deepmd/tf2/model/ener_model.py b/deepmd/tf2/model/ener_model.py index ca019bbaaa..00d2e22004 100644 --- a/deepmd/tf2/model/ener_model.py +++ b/deepmd/tf2/model/ener_model.py @@ -11,6 +11,8 @@ ) +@BaseModel.register("sezm_ener") +@BaseModel.register("dpa4_ener") @BaseModel.register("ener") class EnergyModel(make_tf2_dp_model_from_dpmodel(EnergyModelDP, DPAtomicModelEnergy)): pass diff --git a/deepmd/tf2/model/model.py b/deepmd/tf2/model/model.py index 8977c3357b..46d1848621 100644 --- a/deepmd/tf2/model/model.py +++ b/deepmd/tf2/model/model.py @@ -107,6 +107,53 @@ def get_zbl_model(data: dict) -> DPZBLModel: ) +def get_sezm_model(data: dict) -> BaseModel: + """Build a DPA4/SeZM energy model from the pt-style model config.""" + data = deepcopy(data) + if "spin" in data: + raise NotImplementedError("Spin DPA4/SeZM models are not supported in TF2.") + if str(data.get("bridging_method", "none")).lower() != "none": + raise NotImplementedError("DPA4/SeZM bridging is not supported in TF2.") + if data.get("lora") is not None: + raise NotImplementedError("DPA4/SeZM LoRA is not supported in TF2.") + if data.get("use_compile"): + raise NotImplementedError("model.use_compile is not supported in TF2.") + if data.get("preset_out_bias"): + raise NotImplementedError("DPA4/SeZM preset_out_bias is not supported in TF2.") + + data.pop("type", None) + data.setdefault("descriptor", {}) + data.setdefault("fitting_net", {}) + data["descriptor"].setdefault("type", "dpa4") + data["fitting_net"].setdefault("type", "dpa4_ener") + if data["descriptor"]["type"] not in ("dpa4", "DPA4", "sezm", "SeZM"): + raise ValueError( + "Model type 'dpa4' requires a DPA4/SeZM descriptor, but got " + f"descriptor type '{data['descriptor']['type']}'." + ) + if data["fitting_net"]["type"] not in ("dpa4_ener", "sezm_ener"): + raise ValueError( + "Model type 'dpa4' requires the DPA4/SeZM energy fitting net, but got " + f"fitting_net type '{data['fitting_net']['type']}'." + ) + + descriptor_exclude_types = [ + list(pair) for pair in (data["descriptor"].get("exclude_types") or []) + ] + if "pair_exclude_types" in data: + pair_exclude_types = [list(pair) for pair in (data["pair_exclude_types"] or [])] + if descriptor_exclude_types and descriptor_exclude_types != pair_exclude_types: + raise ValueError( + "DPA4/SeZM pair_exclude_types and descriptor.exclude_types must " + "match when both are provided." + ) + else: + pair_exclude_types = descriptor_exclude_types + data["pair_exclude_types"] = pair_exclude_types + data["descriptor"]["exclude_types"] = deepcopy(pair_exclude_types) + return get_standard_model(data) + + def get_model(data: dict) -> BaseModel: """Get a model from a dictionary. @@ -123,5 +170,7 @@ def get_model(data: dict) -> BaseModel: return get_zbl_model(data) else: return get_standard_model(data) + elif model_type in ("SeZM", "sezm", "DPA4", "dpa4"): + return get_sezm_model(data) else: return BaseModel.get_class_by_type(model_type).get_model(data) diff --git a/deepmd/tf2/train/trainer.py b/deepmd/tf2/train/trainer.py index ada601b077..448051462e 100644 --- a/deepmd/tf2/train/trainer.py +++ b/deepmd/tf2/train/trainer.py @@ -1415,7 +1415,7 @@ def _translate_model_ret_to_loss_dict( None, ) if not callable(translated_output_def): - return model_ret + return self._match_label_shapes(model_ret, label_dict) output_defs = translated_output_def() model_pred = {} for output_key, output_def in output_defs.items(): @@ -1434,7 +1434,36 @@ def _translate_model_ret_to_loss_dict( and "virial" not in model_pred ): model_pred["virial"] = label_dict["virial"] - return model_pred + return self._match_label_shapes(model_pred, label_dict) + + @classmethod + def _match_label_shapes( + cls, + model_dict: dict[str, Any], + label_dict: dict[str, Any] | None, + ) -> dict[str, Any]: + """Match equivalent flattened model outputs to label tensor shapes.""" + if label_dict is None: + return model_dict + force_hat = model_dict.get("force") + force = label_dict.get("force") + if force_hat is None or force is None: + return model_dict + force_hat_shape = cls._static_shape(force_hat) + force_shape = cls._static_shape(force) + if ( + force_hat_shape is None + or force_shape is None + or force_hat_shape == force_shape + or np.prod(force_hat_shape, dtype=np.int64) + != np.prod(force_shape, dtype=np.int64) + ): + return model_dict + model_dict = dict(model_dict) + model_dict["force"] = to_tensorflow_array( + tf.reshape(to_tf_tensor(force_hat), force_shape) + ) + return model_dict @classmethod def _match_output_rank(cls, value: Any, output_def: Any) -> Any: @@ -1452,6 +1481,23 @@ def _match_output_rank(cls, value: Any, output_def: Any) -> Any: else: value = squeeze(axis) + @staticmethod + def _static_shape(value: Any) -> tuple[int, ...] | None: + shape = getattr(value, "shape", None) + if shape is None: + return None + dims = [] + try: + iterator = iter(shape) + except TypeError: + return None + for dim in iterator: + dim = getattr(dim, "value", dim) + if not isinstance(dim, int): + return None + dims.append(dim) + return tuple(dims) + @staticmethod def _shape_rank(value: Any) -> int | None: shape = getattr(value, "shape", None) diff --git a/source/tests/consistent/descriptor/test_dpa4.py b/source/tests/consistent/descriptor/test_dpa4.py index e6f3216bd4..166f0c7498 100644 --- a/source/tests/consistent/descriptor/test_dpa4.py +++ b/source/tests/consistent/descriptor/test_dpa4.py @@ -21,6 +21,7 @@ from ..common import ( INSTALLED_PT, INSTALLED_PT_EXPT, + INSTALLED_TF2, CommonTest, parameterized_cases, ) @@ -36,6 +37,10 @@ from deepmd.pt_expt.descriptor.dpa4 import DescrptDPA4 as DescrptDPA4PTExpt else: DescrptDPA4PTExpt = None +if INSTALLED_TF2: + from deepmd.tf2.descriptor.dpa4 import DescrptDPA4 as DescrptDPA4TF2 +else: + DescrptDPA4TF2 = None # not implemented DescrptDPA4TF = None @@ -150,12 +155,14 @@ def skip_pt(self) -> bool: skip_dp = False skip_tf = True + skip_tf2 = not INSTALLED_TF2 or DescrptDPA4TF2 is None skip_jax = True skip_pd = True skip_pt_expt = not INSTALLED_PT_EXPT skip_array_api_strict = True tf_class = DescrptDPA4TF + tf2_class = DescrptDPA4TF2 dp_class = DescrptDPA4DP pt_class = DescrptDPA4PT pt_expt_class = DescrptDPA4PTExpt @@ -234,6 +241,16 @@ def eval_pt_expt(self, pt_expt_obj: Any) -> Any: mixed_types=True, ) + def eval_tf2(self, tf2_obj: Any) -> Any: + return self.eval_tf2_descriptor( + tf2_obj, + self.natoms, + self.coords, + self.atype, + self.box, + mixed_types=True, + ) + def extract_ret(self, ret: Any, backend) -> tuple[np.ndarray, ...]: return (ret[0],) diff --git a/source/tests/consistent/fitting/test_dpa4_ener.py b/source/tests/consistent/fitting/test_dpa4_ener.py index ee10b54343..1e9055fe36 100644 --- a/source/tests/consistent/fitting/test_dpa4_ener.py +++ b/source/tests/consistent/fitting/test_dpa4_ener.py @@ -6,6 +6,9 @@ import numpy as np +from deepmd.dpmodel.common import ( + to_numpy_array, +) from deepmd.dpmodel.fitting.dpa4_ener import SeZMEnergyFittingNet as SeZMEnerFittingDP from deepmd.env import ( GLOBAL_NP_FLOAT_PRECISION, @@ -17,6 +20,7 @@ from ..common import ( INSTALLED_PT, INSTALLED_PT_EXPT, + INSTALLED_TF2, CommonTest, parameterized_cases, ) @@ -40,6 +44,13 @@ from deepmd.pt_expt.utils.env import DEVICE as PT_EXPT_DEVICE else: SeZMEnerFittingPTExpt = None +if INSTALLED_TF2: + from deepmd.tf2.common import ( + to_tensorflow_array, + ) + from deepmd.tf2.fitting.dpa4_ener import SeZMEnergyFittingNet as SeZMEnerFittingTF2 +else: + SeZMEnerFittingTF2 = None # not implemented SeZMEnerFittingTF = None @@ -74,12 +85,14 @@ def skip_pt(self) -> bool: skip_dp = False skip_tf = True + skip_tf2 = not INSTALLED_TF2 or SeZMEnerFittingTF2 is None skip_jax = True skip_pd = True skip_pt_expt = not INSTALLED_PT_EXPT skip_array_api_strict = True tf_class = SeZMEnerFittingTF + tf2_class = SeZMEnerFittingTF2 dp_class = SeZMEnerFittingDP pt_class = SeZMEnerFittingPT pt_expt_class = SeZMEnerFittingPTExpt @@ -138,6 +151,14 @@ def eval_pt_expt(self, pt_expt_obj: Any) -> Any: .numpy() ) + def eval_tf2(self, tf2_obj: Any) -> Any: + return to_numpy_array( + tf2_obj( + to_tensorflow_array(self.inputs), + to_tensorflow_array(self.atype.reshape(1, -1)), + )["energy"] + ) + def extract_ret(self, ret: Any, backend) -> tuple[np.ndarray, ...]: return (ret,) From 789aa79ca4e5ab702de12e1f372005114b5cc0df Mon Sep 17 00:00:00 2001 From: Jinzhe Zeng Date: Wed, 8 Jul 2026 21:21:09 +0800 Subject: [PATCH 2/4] fix(tf2): address dpa4 ci feedback --- deepmd/dpmodel/array_api.py | 16 +++- deepmd/tf2/common.py | 20 +++-- deepmd/tf2/descriptor/se_atten_v2.py | 6 +- deepmd/tf2/train/trainer.py | 2 +- .../tests/consistent/descriptor/test_dpa4.py | 55 +++++++++++--- .../consistent/fitting/test_dpa4_ener.py | 76 ++++++++++++++----- 6 files changed, 134 insertions(+), 41 deletions(-) diff --git a/deepmd/dpmodel/array_api.py b/deepmd/dpmodel/array_api.py index 4e2ace9bbe..5b4416e54e 100644 --- a/deepmd/dpmodel/array_api.py +++ b/deepmd/dpmodel/array_api.py @@ -293,7 +293,21 @@ def xp_maximum_at(x: Array, indices: Array, values: Array) -> Array: indices_tensor, tf.shape(x_tensor, out_type=tf.int64)[0], ) - return xp.asarray(tf.maximum(x_tensor, reduced)) + segment_counts = tf.math.unsorted_segment_sum( + tf.ones_like(indices_tensor, dtype=tf.int32), + indices_tensor, + tf.shape(x_tensor, out_type=tf.int64)[0], + ) + touched = segment_counts > 0 + touched_shape = tf.concat( + [ + tf.reshape(tf.shape(x_tensor, out_type=tf.int64)[0], (1,)), + tf.ones(tf.rank(x_tensor) - 1, dtype=tf.int64), + ], + axis=0, + ) + touched = tf.reshape(touched, touched_shape) + return xp.asarray(tf.where(touched, tf.maximum(x_tensor, reduced), x_tensor)) else: # Fallback for array_api_strict: basic indexing only. n = indices.shape[0] diff --git a/deepmd/tf2/common.py b/deepmd/tf2/common.py index f3dfc8b12a..9e95aa44a0 100644 --- a/deepmd/tf2/common.py +++ b/deepmd/tf2/common.py @@ -361,10 +361,12 @@ def visit(value: Any) -> None: def __getattribute__(self, name: str) -> Any: if not name.startswith("_tf2_"): - array_attrs = object.__getattribute__( - self, - "_tf2_array_variable_attr_names", - )() + try: + array_attrs = object.__getattribute__( + self, "_tf2_array_variable_attrs" + ) + except AttributeError: + array_attrs = () if name in array_attrs: storage_name = object.__getattribute__( self, @@ -373,10 +375,12 @@ def __getattribute__(self, name: str) -> Any: variable = object.__getattribute__(self, storage_name) return None if variable is None else to_tensorflow_array(variable) - list_attrs = object.__getattribute__( - self, - "_tf2_array_variable_list_attr_names", - )() + try: + list_attrs = object.__getattribute__( + self, "_tf2_array_variable_list_attrs" + ) + except AttributeError: + list_attrs = () if name in list_attrs: storage_name = object.__getattribute__( self, diff --git a/deepmd/tf2/descriptor/se_atten_v2.py b/deepmd/tf2/descriptor/se_atten_v2.py index d343226da9..97fd238a46 100644 --- a/deepmd/tf2/descriptor/se_atten_v2.py +++ b/deepmd/tf2/descriptor/se_atten_v2.py @@ -14,7 +14,11 @@ @BaseDescriptor.register("se_atten_v2") class DescrptSeAttenV2(DescrptDPA1, DescrptSeAttenV2DP): - pass + @classmethod + def deserialize(cls, data: dict) -> "DescrptSeAttenV2": + obj = DescrptSeAttenV2DP.deserialize.__func__(cls, data) + obj._refresh_tf2_trackable_lists() + return obj register_dpmodel_mapping( diff --git a/deepmd/tf2/train/trainer.py b/deepmd/tf2/train/trainer.py index 448051462e..aaca697949 100644 --- a/deepmd/tf2/train/trainer.py +++ b/deepmd/tf2/train/trainer.py @@ -1489,7 +1489,7 @@ def _static_shape(value: Any) -> tuple[int, ...] | None: dims = [] try: iterator = iter(shape) - except TypeError: + except (TypeError, ValueError): return None for dim in iterator: dim = getattr(dim, "value", dim) diff --git a/source/tests/consistent/descriptor/test_dpa4.py b/source/tests/consistent/descriptor/test_dpa4.py index 166f0c7498..f9151209d0 100644 --- a/source/tests/consistent/descriptor/test_dpa4.py +++ b/source/tests/consistent/descriptor/test_dpa4.py @@ -29,14 +29,36 @@ DescriptorTest, ) -if INSTALLED_PT: - from deepmd.pt.model.descriptor.sezm import DescrptSeZM as DescrptDPA4PT -else: - DescrptDPA4PT = None -if INSTALLED_PT_EXPT: - from deepmd.pt_expt.descriptor.dpa4 import DescrptDPA4 as DescrptDPA4PTExpt -else: - DescrptDPA4PTExpt = None +DescrptDPA4PT = None +DescrptDPA4PTExpt = None + + +def _get_descrpt_dpa4_pt() -> type | None: + global DescrptDPA4PT + if not INSTALLED_PT: + return None + if DescrptDPA4PT is None: + from deepmd.pt.model.descriptor.sezm import ( + DescrptSeZM, + ) + + DescrptDPA4PT = DescrptSeZM + return DescrptDPA4PT + + +def _get_descrpt_dpa4_pt_expt() -> type | None: + global DescrptDPA4PTExpt + if not INSTALLED_PT_EXPT: + return None + if DescrptDPA4PTExpt is None: + from deepmd.pt_expt.descriptor.dpa4 import ( + DescrptDPA4, + ) + + DescrptDPA4PTExpt = DescrptDPA4 + return DescrptDPA4PTExpt + + if INSTALLED_TF2: from deepmd.tf2.descriptor.dpa4 import DescrptDPA4 as DescrptDPA4TF2 else: @@ -151,21 +173,30 @@ def data(self) -> dict: @property def skip_pt(self) -> bool: - return CommonTest.skip_pt + return CommonTest.skip_pt or _get_descrpt_dpa4_pt() is None + + @property + def skip_pt_expt(self) -> bool: + return not INSTALLED_PT_EXPT or _get_descrpt_dpa4_pt_expt() is None + + @property + def pt_class(self) -> type | None: + return _get_descrpt_dpa4_pt() + + @property + def pt_expt_class(self) -> type | None: + return _get_descrpt_dpa4_pt_expt() skip_dp = False skip_tf = True skip_tf2 = not INSTALLED_TF2 or DescrptDPA4TF2 is None skip_jax = True skip_pd = True - skip_pt_expt = not INSTALLED_PT_EXPT skip_array_api_strict = True tf_class = DescrptDPA4TF tf2_class = DescrptDPA4TF2 dp_class = DescrptDPA4DP - pt_class = DescrptDPA4PT - pt_expt_class = DescrptDPA4PTExpt jax_class = None pd_class = None array_api_strict_class = None diff --git a/source/tests/consistent/fitting/test_dpa4_ener.py b/source/tests/consistent/fitting/test_dpa4_ener.py index 1e9055fe36..f85451803d 100644 --- a/source/tests/consistent/fitting/test_dpa4_ener.py +++ b/source/tests/consistent/fitting/test_dpa4_ener.py @@ -28,22 +28,53 @@ FittingTest, ) -if INSTALLED_PT: - import torch +torch = None +PT_DEVICE = None +PT_EXPT_DEVICE = None +SeZMEnerFittingPT = None +SeZMEnerFittingPTExpt = None + + +def _get_sezm_ener_fitting_pt() -> type | None: + global PT_DEVICE, SeZMEnerFittingPT, torch + if not INSTALLED_PT: + return None + if SeZMEnerFittingPT is None: + import torch as torch_module + + from deepmd.pt.model.task.sezm_ener import ( + SeZMEnergyFittingNet, + ) + from deepmd.pt.utils.env import ( + DEVICE, + ) + + torch = torch_module + PT_DEVICE = DEVICE + SeZMEnerFittingPT = SeZMEnergyFittingNet + return SeZMEnerFittingPT + + +def _get_sezm_ener_fitting_pt_expt() -> type | None: + global PT_EXPT_DEVICE, SeZMEnerFittingPTExpt, torch + if not INSTALLED_PT_EXPT: + return None + if SeZMEnerFittingPTExpt is None: + import torch as torch_module + + from deepmd.pt_expt.fitting.dpa4_ener import ( + SeZMEnergyFittingNet, + ) + from deepmd.pt_expt.utils.env import ( + DEVICE, + ) + + torch = torch_module + PT_EXPT_DEVICE = DEVICE + SeZMEnerFittingPTExpt = SeZMEnergyFittingNet + return SeZMEnerFittingPTExpt - from deepmd.pt.model.task.sezm_ener import SeZMEnergyFittingNet as SeZMEnerFittingPT - from deepmd.pt.utils.env import DEVICE as PT_DEVICE -else: - SeZMEnerFittingPT = None -if INSTALLED_PT_EXPT: - import torch - from deepmd.pt_expt.fitting.dpa4_ener import ( - SeZMEnergyFittingNet as SeZMEnerFittingPTExpt, - ) - from deepmd.pt_expt.utils.env import DEVICE as PT_EXPT_DEVICE -else: - SeZMEnerFittingPTExpt = None if INSTALLED_TF2: from deepmd.tf2.common import ( to_tensorflow_array, @@ -81,21 +112,30 @@ def data(self) -> dict: @property def skip_pt(self) -> bool: - return CommonTest.skip_pt + return CommonTest.skip_pt or _get_sezm_ener_fitting_pt() is None + + @property + def skip_pt_expt(self) -> bool: + return not INSTALLED_PT_EXPT or _get_sezm_ener_fitting_pt_expt() is None + + @property + def pt_class(self) -> type | None: + return _get_sezm_ener_fitting_pt() + + @property + def pt_expt_class(self) -> type | None: + return _get_sezm_ener_fitting_pt_expt() skip_dp = False skip_tf = True skip_tf2 = not INSTALLED_TF2 or SeZMEnerFittingTF2 is None skip_jax = True skip_pd = True - skip_pt_expt = not INSTALLED_PT_EXPT skip_array_api_strict = True tf_class = SeZMEnerFittingTF tf2_class = SeZMEnerFittingTF2 dp_class = SeZMEnerFittingDP - pt_class = SeZMEnerFittingPT - pt_expt_class = SeZMEnerFittingPTExpt jax_class = None pd_class = None array_api_strict_class = None From 8136fb11b0fdb89042abbe6c166c4be387b598b8 Mon Sep 17 00:00:00 2001 From: Jinzhe Zeng Date: Wed, 8 Jul 2026 21:28:29 +0800 Subject: [PATCH 3/4] test(tf2): revert dpa4 consistent test churn --- .../tests/consistent/descriptor/test_dpa4.py | 55 +++----------- .../consistent/fitting/test_dpa4_ener.py | 76 +++++-------------- 2 files changed, 30 insertions(+), 101 deletions(-) diff --git a/source/tests/consistent/descriptor/test_dpa4.py b/source/tests/consistent/descriptor/test_dpa4.py index f9151209d0..166f0c7498 100644 --- a/source/tests/consistent/descriptor/test_dpa4.py +++ b/source/tests/consistent/descriptor/test_dpa4.py @@ -29,36 +29,14 @@ DescriptorTest, ) -DescrptDPA4PT = None -DescrptDPA4PTExpt = None - - -def _get_descrpt_dpa4_pt() -> type | None: - global DescrptDPA4PT - if not INSTALLED_PT: - return None - if DescrptDPA4PT is None: - from deepmd.pt.model.descriptor.sezm import ( - DescrptSeZM, - ) - - DescrptDPA4PT = DescrptSeZM - return DescrptDPA4PT - - -def _get_descrpt_dpa4_pt_expt() -> type | None: - global DescrptDPA4PTExpt - if not INSTALLED_PT_EXPT: - return None - if DescrptDPA4PTExpt is None: - from deepmd.pt_expt.descriptor.dpa4 import ( - DescrptDPA4, - ) - - DescrptDPA4PTExpt = DescrptDPA4 - return DescrptDPA4PTExpt - - +if INSTALLED_PT: + from deepmd.pt.model.descriptor.sezm import DescrptSeZM as DescrptDPA4PT +else: + DescrptDPA4PT = None +if INSTALLED_PT_EXPT: + from deepmd.pt_expt.descriptor.dpa4 import DescrptDPA4 as DescrptDPA4PTExpt +else: + DescrptDPA4PTExpt = None if INSTALLED_TF2: from deepmd.tf2.descriptor.dpa4 import DescrptDPA4 as DescrptDPA4TF2 else: @@ -173,30 +151,21 @@ def data(self) -> dict: @property def skip_pt(self) -> bool: - return CommonTest.skip_pt or _get_descrpt_dpa4_pt() is None - - @property - def skip_pt_expt(self) -> bool: - return not INSTALLED_PT_EXPT or _get_descrpt_dpa4_pt_expt() is None - - @property - def pt_class(self) -> type | None: - return _get_descrpt_dpa4_pt() - - @property - def pt_expt_class(self) -> type | None: - return _get_descrpt_dpa4_pt_expt() + return CommonTest.skip_pt skip_dp = False skip_tf = True skip_tf2 = not INSTALLED_TF2 or DescrptDPA4TF2 is None skip_jax = True skip_pd = True + skip_pt_expt = not INSTALLED_PT_EXPT skip_array_api_strict = True tf_class = DescrptDPA4TF tf2_class = DescrptDPA4TF2 dp_class = DescrptDPA4DP + pt_class = DescrptDPA4PT + pt_expt_class = DescrptDPA4PTExpt jax_class = None pd_class = None array_api_strict_class = None diff --git a/source/tests/consistent/fitting/test_dpa4_ener.py b/source/tests/consistent/fitting/test_dpa4_ener.py index f85451803d..1e9055fe36 100644 --- a/source/tests/consistent/fitting/test_dpa4_ener.py +++ b/source/tests/consistent/fitting/test_dpa4_ener.py @@ -28,53 +28,22 @@ FittingTest, ) -torch = None -PT_DEVICE = None -PT_EXPT_DEVICE = None -SeZMEnerFittingPT = None -SeZMEnerFittingPTExpt = None - - -def _get_sezm_ener_fitting_pt() -> type | None: - global PT_DEVICE, SeZMEnerFittingPT, torch - if not INSTALLED_PT: - return None - if SeZMEnerFittingPT is None: - import torch as torch_module - - from deepmd.pt.model.task.sezm_ener import ( - SeZMEnergyFittingNet, - ) - from deepmd.pt.utils.env import ( - DEVICE, - ) - - torch = torch_module - PT_DEVICE = DEVICE - SeZMEnerFittingPT = SeZMEnergyFittingNet - return SeZMEnerFittingPT - - -def _get_sezm_ener_fitting_pt_expt() -> type | None: - global PT_EXPT_DEVICE, SeZMEnerFittingPTExpt, torch - if not INSTALLED_PT_EXPT: - return None - if SeZMEnerFittingPTExpt is None: - import torch as torch_module - - from deepmd.pt_expt.fitting.dpa4_ener import ( - SeZMEnergyFittingNet, - ) - from deepmd.pt_expt.utils.env import ( - DEVICE, - ) - - torch = torch_module - PT_EXPT_DEVICE = DEVICE - SeZMEnerFittingPTExpt = SeZMEnergyFittingNet - return SeZMEnerFittingPTExpt +if INSTALLED_PT: + import torch + from deepmd.pt.model.task.sezm_ener import SeZMEnergyFittingNet as SeZMEnerFittingPT + from deepmd.pt.utils.env import DEVICE as PT_DEVICE +else: + SeZMEnerFittingPT = None +if INSTALLED_PT_EXPT: + import torch + from deepmd.pt_expt.fitting.dpa4_ener import ( + SeZMEnergyFittingNet as SeZMEnerFittingPTExpt, + ) + from deepmd.pt_expt.utils.env import DEVICE as PT_EXPT_DEVICE +else: + SeZMEnerFittingPTExpt = None if INSTALLED_TF2: from deepmd.tf2.common import ( to_tensorflow_array, @@ -112,30 +81,21 @@ def data(self) -> dict: @property def skip_pt(self) -> bool: - return CommonTest.skip_pt or _get_sezm_ener_fitting_pt() is None - - @property - def skip_pt_expt(self) -> bool: - return not INSTALLED_PT_EXPT or _get_sezm_ener_fitting_pt_expt() is None - - @property - def pt_class(self) -> type | None: - return _get_sezm_ener_fitting_pt() - - @property - def pt_expt_class(self) -> type | None: - return _get_sezm_ener_fitting_pt_expt() + return CommonTest.skip_pt skip_dp = False skip_tf = True skip_tf2 = not INSTALLED_TF2 or SeZMEnerFittingTF2 is None skip_jax = True skip_pd = True + skip_pt_expt = not INSTALLED_PT_EXPT skip_array_api_strict = True tf_class = SeZMEnerFittingTF tf2_class = SeZMEnerFittingTF2 dp_class = SeZMEnerFittingDP + pt_class = SeZMEnerFittingPT + pt_expt_class = SeZMEnerFittingPTExpt jax_class = None pd_class = None array_api_strict_class = None From 4328329eb47f43fac1b0b83cf4887b9b308490ba Mon Sep 17 00:00:00 2001 From: njzjz-bot Date: Sat, 11 Jul 2026 14:22:07 +0800 Subject: [PATCH 4/4] fix(tf2): address dpa4 review feedback Coding-Agent: Codex Codex-Version: codex-cli 0.144.1 Model: gpt-5.6-sol Reasoning-Effort: xhigh --- deepmd/tf2/descriptor/dpa4.py | 5 +- deepmd/tf2/descriptor/se_atten_v2.py | 4 +- deepmd/tf2/fitting/dpa4_ener.py | 1 - deepmd/tf2/model/model.py | 4 +- source/tests/tf2/test_dpa4.py | 92 ++++++++++++++++++++++++++ source/tests/tf2/test_model_factory.py | 92 ++++++++++++++++++++++++++ source/tests/tf2/test_training.py | 39 +++++++++++ 7 files changed, 229 insertions(+), 8 deletions(-) create mode 100644 source/tests/tf2/test_dpa4.py create mode 100644 source/tests/tf2/test_model_factory.py diff --git a/deepmd/tf2/descriptor/dpa4.py b/deepmd/tf2/descriptor/dpa4.py index b706cac4f9..fee20dfb0f 100644 --- a/deepmd/tf2/descriptor/dpa4.py +++ b/deepmd/tf2/descriptor/dpa4.py @@ -34,8 +34,6 @@ from deepmd.tf2.descriptor.base_descriptor import ( BaseDescriptor, ) -from deepmd.tf2.utils import exclude_mask as _tf2_exclude_mask # noqa: F401 -from deepmd.tf2.utils import network as _tf2_network # noqa: F401 @tf2_module @@ -111,6 +109,8 @@ class WignerDCalculator(WignerDCalculatorDP): "SO3Linear": ("weight", "bias"), "FocusLinear": ("weight", "bias"), "ChannelLinear": ("weight", "bias"), + "FrameContract": ("weight",), + "FrameExpand": ("weight",), "SO2Linear": ("weight_m0", "bias0"), "DynamicRadialDegreeMixer": ("weight", "channel_basis"), "SO2Convolution": ( @@ -130,6 +130,7 @@ class WignerDCalculator(WignerDCalculatorDP): } _TRAINABLE_LIST_ATTRS: dict[str, tuple[str, ...]] = { + "SeZMInteractionBlock": ("adam_ffn_layer_scales",), "SO2Linear": ("weight_m",), "SO2Convolution": ("adam_so2_layer_scales",), } diff --git a/deepmd/tf2/descriptor/se_atten_v2.py b/deepmd/tf2/descriptor/se_atten_v2.py index 97fd238a46..dbb58aa0b5 100644 --- a/deepmd/tf2/descriptor/se_atten_v2.py +++ b/deepmd/tf2/descriptor/se_atten_v2.py @@ -16,9 +16,7 @@ class DescrptSeAttenV2(DescrptDPA1, DescrptSeAttenV2DP): @classmethod def deserialize(cls, data: dict) -> "DescrptSeAttenV2": - obj = DescrptSeAttenV2DP.deserialize.__func__(cls, data) - obj._refresh_tf2_trackable_lists() - return obj + return DescrptSeAttenV2DP.deserialize.__func__(cls, data) register_dpmodel_mapping( diff --git a/deepmd/tf2/fitting/dpa4_ener.py b/deepmd/tf2/fitting/dpa4_ener.py index d20e06060c..4dff6453a2 100644 --- a/deepmd/tf2/fitting/dpa4_ener.py +++ b/deepmd/tf2/fitting/dpa4_ener.py @@ -18,7 +18,6 @@ from deepmd.tf2.fitting.base_fitting import ( BaseFitting, ) -from deepmd.tf2.utils import network as _tf2_network # noqa: F401 @tf2_module diff --git a/deepmd/tf2/model/model.py b/deepmd/tf2/model/model.py index 46d1848621..5824079014 100644 --- a/deepmd/tf2/model/model.py +++ b/deepmd/tf2/model/model.py @@ -122,8 +122,8 @@ def get_sezm_model(data: dict) -> BaseModel: raise NotImplementedError("DPA4/SeZM preset_out_bias is not supported in TF2.") data.pop("type", None) - data.setdefault("descriptor", {}) - data.setdefault("fitting_net", {}) + data["descriptor"] = data.get("descriptor") or {} + data["fitting_net"] = data.get("fitting_net") or {} data["descriptor"].setdefault("type", "dpa4") data["fitting_net"].setdefault("type", "dpa4_ener") if data["descriptor"]["type"] not in ("dpa4", "DPA4", "sezm", "SeZM"): diff --git a/source/tests/tf2/test_dpa4.py b/source/tests/tf2/test_dpa4.py new file mode 100644 index 0000000000..d49df1dcd3 --- /dev/null +++ b/source/tests/tf2/test_dpa4.py @@ -0,0 +1,92 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Focused tests for TF2 DPA4 trainable and trackable state.""" + +import os + +import numpy as np +import pytest + +if os.environ.get("DP_TEST_TF2_ONLY") != "1": + pytest.skip( + "TF2 tests require DP_TEST_TF2_ONLY=1", + allow_module_level=True, + ) + +from deepmd.tf2.descriptor.dpa4 import ( + DescrptDPA4, + _iter_object_tree, +) +from deepmd.tf2.env import ( + tf, +) + + +def _make_trainable_descriptor() -> DescrptDPA4: + """Build a small descriptor that enables the optional trainable leaves.""" + return DescrptDPA4( + ntypes=2, + sel=4, + rcut=4.0, + channels=4, + n_radial=4, + lmax=1, + mmax=1, + n_blocks=1, + grid_branch=0, + layer_scale=True, + message_node_so3=True, + random_gamma=False, + precision="float64", + trainable=True, + seed=20260711, + ) + + +def _assert_optional_weights_are_tracked(descriptor: DescrptDPA4) -> None: + """Assert optional DPA4 variables are trainable TensorFlow trackables.""" + modules = list(_iter_object_tree(descriptor)) + tracked_ids = {id(variable) for variable in descriptor.trainable_variables} + + frame_modules = [ + module + for module in modules + if type(module).__name__ in {"FrameContract", "FrameExpand"} + ] + assert {type(module).__name__ for module in frame_modules} == { + "FrameContract", + "FrameExpand", + } + for module in frame_modules: + variable = object.__getattribute__(module, "_tf2_weight_variable") + assert isinstance(variable, tf.Variable) + assert variable.trainable + assert id(variable) in tracked_ids + + interaction_blocks = [ + module for module in modules if type(module).__name__ == "SeZMInteractionBlock" + ] + assert interaction_blocks + for block in interaction_blocks: + variables = object.__getattribute__( + block, + "_tf2_adam_ffn_layer_scales_variables", + ) + assert variables + assert all(variable.trainable for variable in variables) + assert all(id(variable) in tracked_ids for variable in variables) + + +def test_optional_dpa4_weights_are_tf2_trainable_variables() -> None: + """Optional cross-grid and FFN LayerScale weights must receive gradients.""" + _assert_optional_weights_are_tracked(_make_trainable_descriptor()) + + +def test_dpa4_deserialize_refreshes_trackable_state() -> None: + """Serialization must preserve values and nested TensorFlow trackables.""" + descriptor = _make_trainable_descriptor() + serialized = descriptor.serialize() + + restored = DescrptDPA4.deserialize(serialized) + + np.testing.assert_equal(restored.serialize(), serialized) + _assert_optional_weights_are_tracked(restored) diff --git a/source/tests/tf2/test_model_factory.py b/source/tests/tf2/test_model_factory.py new file mode 100644 index 0000000000..e8c381a15e --- /dev/null +++ b/source/tests/tf2/test_model_factory.py @@ -0,0 +1,92 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Tests for TF2 DPA4/SeZM model-factory validation.""" + +import os + +import pytest + +if os.environ.get("DP_TEST_TF2_ONLY") != "1": + pytest.skip( + "TF2 tests require DP_TEST_TF2_ONLY=1", + allow_module_level=True, + ) + +from deepmd.tf2.model import model as model_module + + +def _base_sezm_config() -> dict: + """Return the smallest config needed to exercise DPA4 factory routing.""" + return { + "type": "dpa4", + "type_map": ["O", "H"], + "descriptor": {"type": "dpa4"}, + "fitting_net": {"type": "dpa4_ener"}, + } + + +def test_null_blocks_receive_dpa4_defaults(monkeypatch: pytest.MonkeyPatch) -> None: + data = _base_sezm_config() + data["descriptor"] = None + data["fitting_net"] = None + monkeypatch.setattr(model_module, "get_standard_model", lambda value: value) + + normalized = model_module.get_model(data) + + assert normalized["descriptor"]["type"] == "dpa4" + assert normalized["fitting_net"]["type"] == "dpa4_ener" + + +@pytest.mark.parametrize( + ("key", "value"), + ( + ("spin", {}), + ("bridging_method", "linear"), + ("lora", {}), + ("use_compile", True), + ("preset_out_bias", [0.0]), + ), +) +def test_rejects_unsupported_features(key: str, value: object) -> None: + data = _base_sezm_config() + data[key] = value + + with pytest.raises(NotImplementedError): + model_module.get_model(data) + + +@pytest.mark.parametrize( + ("section", "model_type"), + (("descriptor", "se_e2_a"), ("fitting_net", "ener")), +) +def test_rejects_incompatible_component_types( + section: str, + model_type: str, +) -> None: + data = _base_sezm_config() + data[section]["type"] = model_type + + with pytest.raises(ValueError): + model_module.get_model(data) + + +def test_rejects_mismatched_exclude_types() -> None: + data = _base_sezm_config() + data["descriptor"]["exclude_types"] = [[0, 1]] + data["pair_exclude_types"] = [[1, 1]] + + with pytest.raises(ValueError): + model_module.get_model(data) + + +def test_descriptor_exclude_types_feed_standard_model( + monkeypatch: pytest.MonkeyPatch, +) -> None: + data = _base_sezm_config() + data["descriptor"] = {"type": "SeZM", "exclude_types": [[0, 1]]} + data["fitting_net"]["type"] = "sezm_ener" + monkeypatch.setattr(model_module, "get_standard_model", lambda value: value) + + normalized = model_module.get_model(data) + + assert normalized["pair_exclude_types"] == [[0, 1]] + assert normalized["descriptor"]["exclude_types"] == [[0, 1]] diff --git a/source/tests/tf2/test_training.py b/source/tests/tf2/test_training.py index 9d7941ade1..8b07b27f81 100644 --- a/source/tests/tf2/test_training.py +++ b/source/tests/tf2/test_training.py @@ -851,6 +851,45 @@ def test_model_ret_translation_uses_translated_output_def() -> None: } +def test_model_ret_translation_reshapes_equivalent_force_layout() -> None: + """A flattened label shape should reshape an equivalent model force.""" + trainer = object.__new__(Trainer) + trainer.models = {"energy": SimpleNamespace()} + force = tf.reshape(tf.range(6, dtype=tf.float64), (1, 2, 3)) + model_ret = {"force": force} + + translated = Trainer._translate_model_ret_to_loss_dict( + trainer, + "energy", + model_ret, + label_dict={"force": tf.zeros((1, 6), dtype=tf.float64)}, + ) + + assert translated is not model_ret + assert tuple(translated["force"].shape) == (1, 6) + np.testing.assert_array_equal( + to_tf_tensor(translated["force"]).numpy(), [[0, 1, 2, 3, 4, 5]] + ) + + +def test_model_ret_translation_preserves_matching_force_layout() -> None: + """An already matching force tensor should remain untouched.""" + trainer = object.__new__(Trainer) + trainer.models = {"energy": SimpleNamespace()} + force = tf.zeros((1, 2, 3), dtype=tf.float64) + model_ret = {"force": force} + + translated = Trainer._translate_model_ret_to_loss_dict( + trainer, + "energy", + model_ret, + label_dict={"force": tf.ones((1, 2, 3), dtype=tf.float64)}, + ) + + assert translated is model_ret + assert translated["force"] is force + + def test_model_ret_translation_only_uses_label_virial_when_not_requested() -> None: trainer = object.__new__(Trainer) trainer.models = {