diff --git a/deepmd/dpmodel/array_api.py b/deepmd/dpmodel/array_api.py index 115242edfb..5b4416e54e 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,32 @@ 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], + ) + 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/dpmodel/descriptor/dpa4.py b/deepmd/dpmodel/descriptor/dpa4.py index d7b4287199..7937a288f9 100644 --- a/deepmd/dpmodel/descriptor/dpa4.py +++ b/deepmd/dpmodel/descriptor/dpa4.py @@ -1341,13 +1341,14 @@ def call( ] # list of (E, lmax+1, C) # === Step 11. Convert to self.dtype and run blocks === - # The block stage is skipped entirely when there are no interaction - # blocks (zero-block descriptor) or no valid edges, sparing the working - # edge-cache dtype cast that only the blocks consume. + # The block stage is skipped entirely for the zero-block descriptor. + # Array operations in the blocks also support an empty edge axis; avoid + # inspecting that dynamic dimension in Python so TF graphs can retrace + # across different atom counts. x = xp.astype(x, get_xp_precision(xp, self.precision)) # (N, D, 1, C) if force_embedding is not None: x = x + xp.astype(force_embedding, get_xp_precision(xp, self.precision)) - if self.blocks and edge_cache.src.shape[0] > 0: + if self.blocks: edge_cache = edge_cache_to_dtype( edge_cache, get_xp_precision(xp, self.precision) ) diff --git a/deepmd/dpmodel/descriptor/dpa4_nn/so2.py b/deepmd/dpmodel/descriptor/dpa4_nn/so2.py index c0c4c0b262..47212831e0 100644 --- a/deepmd/dpmodel/descriptor/dpa4_nn/so2.py +++ b/deepmd/dpmodel/descriptor/dpa4_nn/so2.py @@ -665,7 +665,7 @@ def _project_radial(self, radial_feat: Array) -> Array: device = array_api_compat.device(radial_feat) radial_m0 = xp.reshape( radial_feat[:, : self.lmax + 1, :], - (radial_feat.shape[0], self.input_dim), + (-1, self.input_dim), ) weight = xp_asarray_nodetach(xp, self.weight[...], device=device) return xp.matmul(radial_m0, weight) @@ -744,9 +744,45 @@ def call(self, x_local: Array, radial_feat: Array) -> Array: Invariant radial/type features with shape (E, D_m, C_wide). """ xp = array_api_compat.array_namespace(x_local) - if x_local.shape != radial_feat.shape: + x_shape = x_local.shape + radial_shape = radial_feat.shape + + def static_rank(shape: Any) -> int | None: + rank = getattr(shape, "rank", None) + if rank is not None: + return int(rank) + try: + return len(shape) + except (TypeError, ValueError): + return None + + def static_dim(shape: Any, axis: int) -> int | None: + try: + dim = shape[axis] + except (IndexError, TypeError, ValueError): + return None + dim = getattr(dim, "value", dim) + return int(dim) if isinstance(dim, (int, np.integer)) else None + + x_rank = static_rank(x_shape) + radial_rank = static_rank(radial_shape) + if (x_rank is not None and x_rank != 3) or ( + radial_rank is not None and radial_rank != 3 + ): + raise ValueError("DynamicRadialDegreeMixer inputs must have rank 3") + if any( + x_dim is not None and radial_dim is not None and x_dim != radial_dim + for x_dim, radial_dim in ( + (static_dim(x_shape, axis), static_dim(radial_shape, axis)) + for axis in range(3) + ) + ): raise ValueError("`x_local` and `radial_feat` must have the same shape") - if x_local.shape[1] != self.reduced_dim or x_local.shape[2] != self.channels: + reduced_dim = static_dim(x_shape, 1) + channel_dim = static_dim(x_shape, 2) + if (reduced_dim is not None and reduced_dim != self.reduced_dim) or ( + channel_dim is not None and channel_dim != self.channels + ): raise ValueError("Input shape is incompatible with this mixer") kernel_flat = self._project_radial(radial_feat) @@ -755,14 +791,10 @@ def call(self, x_local: Array, radial_feat: Array) -> Array: return xp.matmul(kernel, x_local) if self.rank > 0: - compact = xp.reshape( - kernel_flat, (x_local.shape[0], self.degree_kernel_size, self.rank) - ) + compact = xp.reshape(kernel_flat, (-1, self.degree_kernel_size, self.rank)) return self._mix_rank_compact(compact, x_local) - compact = xp.reshape( - kernel_flat, (x_local.shape[0], self.degree_kernel_size, self.channels) - ) + compact = xp.reshape(kernel_flat, (-1, self.degree_kernel_size, self.channels)) kernel = self._scatter_channel_kernel(compact) # einsum("eoic,eic->eoc"): contract l_in i per channel c (no channel mix). return xp.sum(kernel * x_local[:, None, :, :], axis=2) @@ -791,12 +823,12 @@ def _mix_rank_compact(self, compact: Array, x_local: Array) -> Array: # via a single matmul, then weight the rank channels by channel_basis. kernel_or = xp.reshape( xp.permute_dims(kernel, (0, 1, 3, 2)), - (x_local.shape[0], self.reduced_dim * self.rank, self.reduced_dim), + (-1, self.reduced_dim * self.rank, self.reduced_dim), ) mixed = xp.matmul(kernel_or, x_local) mixed = xp.reshape( mixed, - (x_local.shape[0], self.reduced_dim, self.rank, self.channels), + (-1, self.reduced_dim, self.rank, self.channels), ) channel_basis = xp.reshape( xp_asarray_nodetach(xp, self.channel_basis[...], device=device), diff --git a/deepmd/dpmodel/fitting/dpa4_ener.py b/deepmd/dpmodel/fitting/dpa4_ener.py index f08097f032..6f7f970947 100644 --- a/deepmd/dpmodel/fitting/dpa4_ener.py +++ b/deepmd/dpmodel/fitting/dpa4_ener.py @@ -92,11 +92,18 @@ def __init__( ) if neuron is None: neuron = [] - if isinstance(trainable, list): - trainable = all(trainable) self.in_dim = int(in_dim) self.out_dim = int(out_dim) self.neuron = [int(nn_dim) for nn_dim in neuron] + if isinstance(trainable, bool): + self.trainable = [trainable] * (len(self.neuron) + 1) + else: + self.trainable = [bool(flag) for flag in trainable] + if len(self.trainable) != len(self.neuron) + 1: + raise ValueError( + "trainable must contain one flag per hidden layer plus " + "one flag for the output layer" + ) self.activation_function = activation_function self.resnet_dt = bool(resnet_dt) self.precision = precision @@ -123,7 +130,7 @@ def __init__( resnet=False, precision=self.precision, seed=child_seed(seed, layer_idx), - trainable=trainable, + trainable=self.trainable[layer_idx], ) ) dim_in = hidden_dim @@ -139,7 +146,7 @@ def __init__( resnet=False, precision=self.precision, seed=child_seed(seed, len(self.neuron) + int(self.case_film_embd)), - trainable=trainable, + trainable=self.trainable[-1], ) def call_until_last(self, xx: Array) -> Array: @@ -181,6 +188,9 @@ def serialize(self) -> dict[str, Any]: "descriptor_dim": self.descriptor_dim, "dim_case_embd": self.dim_case_embd, "case_film_embd": self.case_film_embd, + # Preserve the effective per-layer freeze policy when backend + # wrappers rebuild this network from its serialized form. + "trainable": self.trainable.copy(), "@variables": variables, } diff --git a/deepmd/pt/model/task/sezm_ener.py b/deepmd/pt/model/task/sezm_ener.py index 0932ec7086..f83fd339f3 100644 --- a/deepmd/pt/model/task/sezm_ener.py +++ b/deepmd/pt/model/task/sezm_ener.py @@ -243,11 +243,18 @@ def __init__( super().__init__() if neuron is None: neuron = [] - if isinstance(trainable, list): - trainable = all(trainable) self.in_dim = int(in_dim) self.out_dim = int(out_dim) self.neuron = [int(nn_dim) for nn_dim in neuron] + if isinstance(trainable, bool): + self.trainable = [trainable] * (len(self.neuron) + 1) + else: + self.trainable = [bool(flag) for flag in trainable] + if len(self.trainable) != len(self.neuron) + 1: + raise ValueError( + "trainable must contain one flag per hidden layer plus " + "one flag for the output layer" + ) self.activation_function = activation_function self.resnet_dt = bool(resnet_dt) self.precision = precision @@ -270,7 +277,7 @@ def __init__( activation_function=self.activation_function, precision=self.precision, seed=child_seed(seed, layer_idx), - trainable=trainable, + trainable=self.trainable[layer_idx], ) ) dim_in = hidden_dim @@ -285,7 +292,7 @@ def __init__( activation_function=self.activation_function, precision=self.precision, seed=child_seed(seed, len(self.neuron)), - trainable=trainable, + trainable=all(self.trainable), ) else: self.case_film = None @@ -300,12 +307,9 @@ def __init__( resnet=False, precision=self.precision, seed=child_seed(seed, len(self.neuron) + int(self.case_film_embd)), - trainable=trainable, + trainable=self.trainable[-1], ) - for param in self.parameters(): - param.requires_grad = trainable - def _apply_input_film( self, xx: torch.Tensor, @@ -395,6 +399,8 @@ def serialize(self) -> dict[str, Any]: "descriptor_dim": self.descriptor_dim, "dim_case_embd": self.dim_case_embd, "case_film_embd": self.case_film_embd, + # Keep the per-layer freeze policy stable across backend round trips. + "trainable": self.trainable.copy(), "@variables": {key: to_numpy_array(value) for key, value in state.items()}, } diff --git a/deepmd/tf2/common.py b/deepmd/tf2/common.py index bb8155a38c..fa0ea960f7 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,68 @@ 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) + trainable_by_name = object.__getattribute__(self, "__dict__").get( + "_tf2_array_variable_trainable", {} + ) + if value is None: + tf.Module.__setattr__(self, storage_name, None) + object.__getattribute__(self, "__dict__").pop(name, None) + return + tensor = to_tf_tensor(value) + variable = tf.Variable( + tensor, + trainable=bool( + trainable_by_name.get( + name, + getattr(self, "trainable", True), + ) + ), + name=name, + ) + tf.Module.__setattr__(self, storage_name, variable) + # The variable-backed accessor owns this value now. Keeping the + # original eager tensor in the public slot doubles parameter RAM. + object.__getattribute__(self, "__dict__").pop(name, None) + + def _set_tf2_array_variable_list(self, name: str, value: Any) -> None: + storage_name = self._tf2_array_variable_list_storage_name(name) + trainable_by_name = object.__getattribute__(self, "__dict__").get( + "_tf2_array_variable_list_trainable", {} + ) + variables = [] + for idx, item in enumerate(value): + tensor = to_tf_tensor(item) + variables.append( + tf.Variable( + tensor, + trainable=bool( + trainable_by_name.get( + name, + getattr(self, "trainable", True), + ) + ), + name=f"{name}_{idx}", + ) + ) + tf.Module.__setattr__(self, storage_name, variables) + object.__getattribute__(self, "__dict__").pop(name, None) + def __init__(self, *args: Any, **kwargs: Any) -> None: tf.Module.__init__(self) super().__init__(*args, **kwargs) @@ -266,11 +331,113 @@ 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_"): + try: + array_attrs = object.__getattribute__( + self, "_tf2_array_variable_attrs" + ) + except AttributeError: + array_attrs = () + 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) + + 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, + "_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..a8a9799e46 --- /dev/null +++ b/deepmd/tf2/descriptor/dpa4.py @@ -0,0 +1,321 @@ +# 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 ( + DynamicRadialDegreeMixer as DynamicRadialDegreeMixerDP, +) +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, +) + + +@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"), + "FrameContract": ("weight",), + "FrameExpand": ("weight",), + "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, ...]] = { + "SeZMInteractionBlock": ("adam_ffn_layer_scales",), + "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_parameter_attr(module: Any, name: str, *, trainable: bool) -> None: + attrs = set(getattr(module, "_tf2_array_variable_attrs", ())) + if name not in attrs: + tf.Module.__setattr__(module, "_tf2_array_variable_attrs", attrs | {name}) + policies = dict(getattr(module, "_tf2_array_variable_trainable", {})) + policies[name] = trainable + tf.Module.__setattr__(module, "_tf2_array_variable_trainable", policies) + + +def _enable_tf2_parameter_list_attr(module: Any, name: str, *, trainable: bool) -> 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}, + ) + policies = dict(getattr(module, "_tf2_array_variable_list_trainable", {})) + policies[name] = trainable + tf.Module.__setattr__(module, "_tf2_array_variable_list_trainable", policies) + + +def _promote_parameters( + module: Any, names: tuple[str, ...], *, trainable: bool +) -> None: + for name in names: + if not hasattr(module, name): + continue + value = getattr(module, name) + if not _is_floating_array(value): + continue + _enable_tf2_parameter_attr(module, name, trainable=trainable) + setattr(module, name, value) + + +def _promote_parameter_lists( + module: Any, names: tuple[str, ...], *, trainable: bool +) -> None: + 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_parameter_list_attr(module, name, trainable=trainable) + setattr(module, name, value) + + +def _promote_trainable_tree(module: Any) -> Any: + root_trainable = bool(getattr(module, "trainable", True)) + for submodule in _iter_object_tree(module): + trainable = root_trainable and bool(getattr(submodule, "trainable", True)) + names = _TRAINABLE_ATTRS.get(type(submodule).__name__) + if names is not None: + _promote_parameters(submodule, names, trainable=trainable) + list_names = _TRAINABLE_LIST_ATTRS.get(type(submodule).__name__) + if list_names is not None: + _promote_parameter_lists(submodule, list_names, trainable=trainable) + return module + + +@tf2_module +class SO2Linear(SO2LinearDP): + def __init__(self, *args: Any, **kwargs: Any) -> None: + super().__init__(*args, **kwargs) + _promote_parameter_lists(self, ("weight_m",), trainable=bool(self.trainable)) + + @classmethod + def deserialize(cls, data: dict) -> "SO2Linear": + obj = super().deserialize(data) + _promote_parameter_lists(obj, ("weight_m",), trainable=bool(obj.trainable)) + return obj + + +register_dpmodel_mapping( + SO2LinearDP, + lambda v: SO2Linear.deserialize(v.serialize()), +) + + +@tf2_module +class DynamicRadialDegreeMixer(DynamicRadialDegreeMixerDP): + """TF2 mixer with runtime shape checks for generalized edge counts.""" + + def call(self, x_local: Any, radial_feat: Any) -> Any: + x_tensor = to_tf_tensor(x_local) + radial_tensor = to_tf_tensor(radial_feat) + assertions = ( + tf.debugging.assert_rank(x_tensor, 3), + tf.debugging.assert_rank(radial_tensor, 3), + tf.debugging.assert_equal( + tf.shape(x_tensor), + tf.shape(radial_tensor), + message="x_local and radial_feat must have the same shape", + ), + tf.debugging.assert_equal(tf.shape(x_tensor)[1], self.reduced_dim), + tf.debugging.assert_equal(tf.shape(x_tensor)[2], self.channels), + ) + with tf.control_dependencies(assertions): + # Runtime assertions establish the contract; ensure_shape carries + # the proven rank into ndtensorflow static metadata. + x_tensor = tf.ensure_shape(tf.identity(x_tensor), [None, None, None]) + radial_tensor = tf.ensure_shape( + tf.identity(radial_tensor), [None, None, None] + ) + return super().call( + xp.asarray(x_tensor), + xp.asarray(radial_tensor), + ) + + +register_dpmodel_mapping( + DynamicRadialDegreeMixerDP, + lambda v: DynamicRadialDegreeMixer.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) + if self.random_gamma: + raise NotImplementedError( + "TF2 DPA4 does not yet support graph-safe random_gamma; " + "set descriptor.random_gamma to false." + ) + _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/descriptor/se_atten_v2.py b/deepmd/tf2/descriptor/se_atten_v2.py index d343226da9..dbb58aa0b5 100644 --- a/deepmd/tf2/descriptor/se_atten_v2.py +++ b/deepmd/tf2/descriptor/se_atten_v2.py @@ -14,7 +14,9 @@ @BaseDescriptor.register("se_atten_v2") class DescrptSeAttenV2(DescrptDPA1, DescrptSeAttenV2DP): - pass + @classmethod + def deserialize(cls, data: dict) -> "DescrptSeAttenV2": + return DescrptSeAttenV2DP.deserialize.__func__(cls, data) register_dpmodel_mapping( 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..4dff6453a2 --- /dev/null +++ b/deepmd/tf2/fitting/dpa4_ener.py @@ -0,0 +1,52 @@ +# 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, +) + + +@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/base_model.py b/deepmd/tf2/model/base_model.py index 9fa57f33c2..fc536187ba 100644 --- a/deepmd/tf2/model/base_model.py +++ b/deepmd/tf2/model/base_model.py @@ -20,8 +20,89 @@ tf, xp, ) +from deepmd.utils.version import ( + check_version_compatibility, +) + + +class BaseModel(make_base_model()): + """TF2 model registry with adapters for regular PT SeZM checkpoints.""" + + _SEZM_MODEL_TYPES = frozenset({"sezm", "dpa4"}) + _SEZM_ATOMIC_TYPES = frozenset({"sezm_atomic"}) -BaseModel = make_base_model() + @classmethod + def deserialize(cls, data: dict[str, Any]) -> "BaseModel": + model_type = str(data.get("type", "standard")).lower() + if model_type in cls._SEZM_MODEL_TYPES: + return cls.deserialize(cls._unwrap_pt_sezm_model(data)) + if model_type in cls._SEZM_ATOMIC_TYPES: + return cls.deserialize(cls._normalize_pt_sezm_atomic(data)) + return super().deserialize(data) + + @staticmethod + def _unwrap_pt_sezm_model(data: dict[str, Any]) -> dict[str, Any]: + """Unwrap PT's model-level SeZM schema after validating its extras.""" + check_version_compatibility(int(data.get("@version", 1)), 1, 1) + if str(data.get("bridging_method", "none")).lower() not in ("none", ""): + raise NotImplementedError( + "PT SeZM/DPA4 checkpoints with bridging are not supported in TF2." + ) + if data.get("lora") is not None: + raise NotImplementedError( + "PT SeZM/DPA4 checkpoints with LoRA are not supported in TF2." + ) + atomic_model = data.get("atomic_model") + if atomic_model is None: + raise ValueError("SeZM/DPA4 model data is missing 'atomic_model'.") + return atomic_model + + @staticmethod + def _normalize_pt_sezm_atomic(data: dict[str, Any]) -> dict[str, Any]: + """Convert PT's energy-only ``sezm_atomic`` schema to ``standard``.""" + data = data.copy() + check_version_compatibility(int(data.get("@version", 2)), 3, 2) + if data.pop("dens_fitting", None) is not None: + raise NotImplementedError( + "PT SeZM/DPA4 checkpoints with a dens head are not supported in TF2." + ) + active_mode = data.pop("active_mode", None) + if active_mode not in (None, "ener"): + raise NotImplementedError( + f"PT SeZM/DPA4 active_mode {active_mode!r} is not supported in TF2." + ) + variables = data.get("@variables") + if isinstance(variables, dict): + data["@variables"] = { + key: value + for key, value in variables.items() + if key in ("out_bias", "out_std") + } + descriptor = data.get("descriptor") + descriptor_config = ( + descriptor.get("config") if isinstance(descriptor, dict) else None + ) + if isinstance(descriptor, dict) and ( + descriptor.get("random_gamma") + or ( + isinstance(descriptor_config, dict) + and descriptor_config.get("random_gamma") + ) + ): + # PT checkpoints may keep the training augmentation enabled, while + # TF2 conversion produces an inference SavedModel. Inference fixes + # the local-Z roll, so normalize this field before construction. + descriptor = descriptor.copy() + if isinstance(descriptor_config, dict): + descriptor_config = descriptor_config.copy() + descriptor_config["random_gamma"] = False + descriptor["config"] = descriptor_config + else: + descriptor["random_gamma"] = False + data["descriptor"] = descriptor + data["@version"] = 2 + data["type"] = "standard" + return data def _collect_model_predict( 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..aee5ff1c94 100644 --- a/deepmd/tf2/model/model.py +++ b/deepmd/tf2/model/model.py @@ -107,6 +107,57 @@ def get_zbl_model(data: dict) -> DPZBLModel: ) +def get_sezm_model(data: dict) -> BaseModel: + """Build a DPA4/SeZM invariant 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["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"): + 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", + "property", + ): + raise ValueError( + "Model type 'dpa4' requires a DPA4/SeZM energy or property 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 []) + ] + pair_exclude_types = [list(pair) for pair in (data.get("pair_exclude_types") or [])] + if pair_exclude_types: + 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 +174,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 8cfa12cbda..70850f2e88 100644 --- a/deepmd/tf2/train/trainer.py +++ b/deepmd/tf2/train/trainer.py @@ -1423,7 +1423,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(): @@ -1442,7 +1442,45 @@ 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 == force_shape and force_hat_shape is not None: + return model_dict + if ( + force_hat_shape is not None + and force_shape is not None + and np.prod(force_hat_shape, dtype=np.int64) + != np.prod(force_shape, dtype=np.int64) + ): + return model_dict + force_hat_tensor = to_tf_tensor(force_hat) + force_tensor = to_tf_tensor(force) + size_assertion = tf.debugging.assert_equal( + tf.size(force_hat_tensor), + tf.size(force_tensor), + message="model and label force tensors must contain the same values", + ) + with tf.control_dependencies([size_assertion]): + reshaped_force = tf.reshape(force_hat_tensor, tf.shape(force_tensor)) + reshaped_force.set_shape(force_tensor.shape) + model_dict = dict(model_dict) + model_dict["force"] = to_tensorflow_array(reshaped_force) + return model_dict @classmethod def _match_output_rank(cls, value: Any, output_def: Any) -> Any: @@ -1460,6 +1498,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, ValueError): + 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 1cfbf2f9a1..23fb7293b9 100644 --- a/source/tests/consistent/descriptor/test_dpa4.py +++ b/source/tests/consistent/descriptor/test_dpa4.py @@ -22,6 +22,7 @@ INSTALLED_ARRAY_API_STRICT, INSTALLED_PT, INSTALLED_PT_EXPT, + INSTALLED_TF2, CommonTest, parameterized_cases, ) @@ -37,6 +38,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 if INSTALLED_ARRAY_API_STRICT: from ...array_api_strict.descriptor.dpa4 import DescrptDPA4 as DescrptDPA4Strict else: @@ -155,12 +160,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 = not INSTALLED_ARRAY_API_STRICT tf_class = DescrptDPA4TF + tf2_class = DescrptDPA4TF2 dp_class = DescrptDPA4DP pt_class = DescrptDPA4PT pt_expt_class = DescrptDPA4PTExpt @@ -239,6 +246,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 eval_array_api_strict(self, array_api_strict_obj: Any) -> Any: return self.eval_array_api_strict_descriptor( array_api_strict_obj, diff --git a/source/tests/consistent/fitting/test_dpa4_ener.py b/source/tests/consistent/fitting/test_dpa4_ener.py index 3d007a9959..e2d80cd975 100644 --- a/source/tests/consistent/fitting/test_dpa4_ener.py +++ b/source/tests/consistent/fitting/test_dpa4_ener.py @@ -21,6 +21,7 @@ INSTALLED_ARRAY_API_STRICT, INSTALLED_PT, INSTALLED_PT_EXPT, + INSTALLED_TF2, CommonTest, parameterized_cases, ) @@ -44,6 +45,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 if INSTALLED_ARRAY_API_STRICT: import array_api_strict @@ -86,12 +94,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 = not INSTALLED_ARRAY_API_STRICT tf_class = SeZMEnerFittingTF + tf2_class = SeZMEnerFittingTF2 dp_class = SeZMEnerFittingDP pt_class = SeZMEnerFittingPT pt_expt_class = SeZMEnerFittingPTExpt @@ -150,6 +160,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 eval_array_api_strict(self, array_api_strict_obj: Any) -> Any: return to_numpy_array( array_api_strict_obj( diff --git a/source/tests/tf2/test_dpa4.py b/source/tests/tf2/test_dpa4.py new file mode 100644 index 0000000000..202ceb56dc --- /dev/null +++ b/source/tests/tf2/test_dpa4.py @@ -0,0 +1,239 @@ +# 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.common import ( + to_tf_tensor, + wrap_tensor, +) +from deepmd.tf2.descriptor.dpa4 import ( + DescrptDPA4, + DynamicRadialDegreeMixer, + _iter_object_tree, +) +from deepmd.tf2.env import ( + tf, +) +from deepmd.tf2.fitting.dpa4_ener import ( + SeZMEnergyFittingNet, +) + + +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) + + +def _make_frozen_descriptor(seed: int) -> DescrptDPA4: + """Build a frozen descriptor whose complete state must remain trackable.""" + return DescrptDPA4( + ntypes=2, + sel=4, + rcut=4.0, + channels=4, + n_radial=4, + lmax=1, + mmax=1, + n_blocks=1, + grid_branch=0, + random_gamma=False, + precision="float64", + trainable=False, + seed=seed, + ) + + +def test_frozen_descriptor_tracks_and_restores_every_parameter(tmp_path) -> None: + """Frozen leaves are non-trainable variables included in checkpoints.""" + source = _make_frozen_descriptor(20260712) + target = _make_frozen_descriptor(20260713) + source_embedding = object.__getattribute__( + source.type_embedding, "_tf2_adam_type_embedding_variable" + ) + target_embedding = object.__getattribute__( + target.type_embedding, "_tf2_adam_type_embedding_variable" + ) + assert not np.array_equal(source_embedding.numpy(), target_embedding.numpy()) + assert source.variables + assert not source.trainable_variables + + checkpoint_path = tf.train.Checkpoint(descriptor=source).save( + str(tmp_path / "descriptor") + ) + tf.train.Checkpoint(descriptor=target).restore(checkpoint_path).assert_consumed() + + np.testing.assert_array_equal(target_embedding.numpy(), source_embedding.numpy()) + + +def test_promoted_parameters_release_public_tensor_shadows() -> None: + """Variable-backed attributes must not retain their original eager tensors.""" + descriptor = _make_trainable_descriptor() + for module in _iter_object_tree(descriptor): + raw_attrs = object.__getattribute__(module, "__dict__") + for name in getattr(module, "_tf2_array_variable_attrs", ()): + assert name not in raw_attrs + for name in getattr(module, "_tf2_array_variable_list_attrs", ()): + assert name not in raw_attrs + + +@pytest.mark.parametrize( + ("policy", "expected_trainable"), + ((False, [False, False, False]), ([False, True], [False, False, True])), +) +def test_fitting_trainability_survives_conversion_and_optimizer_step( + policy: bool | list[bool], + expected_trainable: list[bool], +) -> None: + """Frozen layers stay fixed while an enabled output layer updates.""" + fitting = SeZMEnergyFittingNet( + ntypes=2, + dim_descrpt=4, + neuron=[4], + trainable=policy, + precision="float64", + mixed_types=True, + seed=20260712, + ) + restored = SeZMEnergyFittingNet.deserialize(fitting.serialize()) + assert [variable.trainable for variable in restored.variables] == expected_trainable + + before = [variable.numpy().copy() for variable in restored.variables] + with tf.GradientTape() as tape: + result = restored( + wrap_tensor(tf.ones((1, 2, 4), dtype=tf.float64)), + wrap_tensor(tf.zeros((1, 2), dtype=tf.int32)), + ) + loss = tf.reduce_sum(to_tf_tensor(result["energy"])) + gradients = tape.gradient(loss, restored.trainable_variables) + if restored.trainable_variables: + tf.keras.optimizers.SGD(0.1).apply_gradients( + zip(gradients, restored.trainable_variables, strict=True) + ) + + changed = [ + not np.array_equal(variable.numpy(), original) + for variable, original in zip(restored.variables, before, strict=True) + ] + assert changed == expected_trainable + + +def test_random_gamma_fails_fast_until_graph_safe_rng_is_supported() -> None: + """TF2 must not silently disable the default random-roll augmentation.""" + with pytest.raises(NotImplementedError, match="random_gamma"): + DescrptDPA4( + ntypes=2, + sel=4, + rcut=4.0, + channels=4, + n_radial=4, + lmax=1, + mmax=1, + n_blocks=1, + precision="float64", + random_gamma=True, + seed=20260712, + ) + + +def test_dynamic_radial_mixer_accepts_unknown_rank_tensor_specs() -> None: + """Runtime rank and shape checks support fully unknown TensorSpecs.""" + mixer = DynamicRadialDegreeMixer( + lmax=1, + mmax=1, + channels=4, + mode="degree_channel", + rank=0, + precision="float64", + seed=20260712, + trainable=True, + ) + + @tf.function( + input_signature=( + tf.TensorSpec(shape=None, dtype=tf.float64), + tf.TensorSpec(shape=None, dtype=tf.float64), + ) + ) + def apply_mixer(x_local: tf.Tensor, radial_feat: tf.Tensor) -> tf.Tensor: + return to_tf_tensor(mixer(wrap_tensor(x_local), wrap_tensor(radial_feat))) + + for nedge in (2, 3): + inputs = tf.ones((nedge, mixer.reduced_dim, mixer.channels), tf.float64) + output = apply_mixer(inputs, inputs) + assert tuple(output.shape) == (nedge, mixer.reduced_dim, mixer.channels) diff --git a/source/tests/tf2/test_dpa4_conversion.py b/source/tests/tf2/test_dpa4_conversion.py new file mode 100644 index 0000000000..7f483bd94a --- /dev/null +++ b/source/tests/tf2/test_dpa4_conversion.py @@ -0,0 +1,116 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""End-to-end PT checkpoint conversion coverage for TF2 DPA4.""" + +import os +import subprocess +import sys +from pathlib import ( + Path, +) + +import pytest + +import deepmd + +if os.environ.get("DP_TEST_TF2_ONLY") != "1": + pytest.skip( + "TF2 tests require DP_TEST_TF2_ONLY=1", + allow_module_level=True, + ) + + +_CONVERSION_SCRIPT = r""" +import copy +import os +import sys +from pathlib import Path + +# Prefer the checked-out source while retaining the installed compiled library. +sys.meta_path[:] = [ + finder + for finder in sys.meta_path + if type(finder).__name__ != "ScikitBuildRedirectingFinder" +] +sys.path.insert(0, os.environ["DEEPMD_SOURCE_ROOT"]) +import deepmd + +compiled_package = os.environ.get("DEEPMD_COMPILED_PACKAGE") +if compiled_package and compiled_package not in deepmd.__path__: + deepmd.__path__.append(compiled_package) + +# Load PT before TensorFlow. Some mixed CUDA plugin builds are not safe when +# Triton is initialized after TensorFlow in the same process. +import torch +from deepmd.pt.model.model import get_model as get_pt_model +from deepmd.pt.train.wrapper import ModelWrapper +from deepmd.pt.utils.serialization import serialize_from_file as serialize_from_pt_file +from deepmd.tf2.env import tf +from deepmd.tf2.utils.serialization import deserialize_to_file as deserialize_to_tf2_file +from deepmd.utils.argcheck import model_args + +output_dir = Path(sys.argv[1]) +model_params = model_args().normalize_value( + { + "type": "dpa4", + "type_map": ["A", "B"], + "descriptor": { + "type": "dpa4", + "sel": 4, + "rcut": 4.0, + "channels": 4, + "n_radial": 4, + "lmax": 1, + "mmax": 1, + "n_blocks": 1, + "radial_so2_mode": "degree_channel", + "precision": "float64", + "seed": 1, + }, + "fitting_net": { + "type": "dpa4_ener", + "neuron": [4], + "precision": "float64", + "seed": 1, + }, + }, + trim_pattern="_.*", +) +pt_model = get_pt_model(copy.deepcopy(model_params)).to(torch.float64) +wrapper = ModelWrapper(pt_model, model_params=copy.deepcopy(model_params)) +pt_path = output_dir / "dpa4.pt" +tf2_path = output_dir / "dpa4.savedmodeltf" +torch.save({"model": wrapper.state_dict()}, pt_path) + +data = serialize_from_pt_file(str(pt_path)) +assert data["model"]["type"] == "SeZM" +deserialize_to_tf2_file(str(tf2_path), data, jit_compile=False) +restored = tf.saved_model.load(str(tf2_path)) +assert callable(restored.call_lower) +assert callable(restored.call_lower_atomic_virial) +""" + + +def test_pt_dpa4_checkpoint_converts_to_tf2_savedmodel(tmp_path) -> None: + """The public ``.pt`` SeZM schema exports through the TF2 adapter.""" + source_root = Path(__file__).parents[3] + compiled_package = next( + ( + package_path + for package_path in deepmd.__path__ + if (Path(package_path) / "lib").is_dir() + ), + "", + ) + env = os.environ.copy() + env["DEEPMD_SOURCE_ROOT"] = str(source_root) + env["DEEPMD_COMPILED_PACKAGE"] = str(compiled_package) + result = subprocess.run( + [sys.executable, "-c", _CONVERSION_SCRIPT, str(tmp_path)], + check=False, + capture_output=True, + text=True, + timeout=180, + env=env, + ) + + assert result.returncode == 0, result.stdout + result.stderr diff --git a/source/tests/tf2/test_model_factory.py b/source/tests/tf2/test_model_factory.py new file mode 100644 index 0000000000..e17bd8a5d4 --- /dev/null +++ b/source/tests/tf2/test_model_factory.py @@ -0,0 +1,148 @@ +# 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 +from deepmd.tf2.model.property_model import ( + PropertyModel, +) +from deepmd.utils.argcheck import ( + model_args, +) + + +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]] + + +def test_normalized_descriptor_exclusions_override_empty_default( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Argcheck's empty model-level default is not an explicit mismatch.""" + data = _base_sezm_config() + data["descriptor"]["exclude_types"] = [[0, 1]] + data = model_args().normalize_value(data, trim_pattern="_.*") + 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]] + + +def test_normalized_dpa4_property_model_is_constructed() -> None: + """The schema-supported invariant property route reaches PropertyModel.""" + data = model_args().normalize_value( + { + "type": "dpa4", + "type_map": ["A", "B"], + "descriptor": { + "type": "dpa4", + "sel": 4, + "rcut": 4.0, + "channels": 4, + "n_radial": 4, + "lmax": 1, + "mmax": 1, + "n_blocks": 1, + "random_gamma": False, + "precision": "float64", + "seed": 1, + }, + "fitting_net": { + "type": "property", + "property_name": "foo", + "task_dim": 3, + "intensive": False, + "neuron": [4], + "precision": "float64", + "seed": 1, + }, + }, + trim_pattern="_.*", + ) + + assert isinstance(model_module.get_model(data), PropertyModel) diff --git a/source/tests/tf2/test_training.py b/source/tests/tf2/test_training.py index dc6324684c..2dbde588d8 100644 --- a/source/tests/tf2/test_training.py +++ b/source/tests/tf2/test_training.py @@ -49,12 +49,18 @@ from deepmd.tf2.model.base_model import ( forward_common_atomic, ) +from deepmd.tf2.model.model import ( + get_model, +) from deepmd.tf2.train.trainer import ( Trainer, ) from deepmd.tf2.utils.jit import ( default_jit_compile, ) +from deepmd.utils.argcheck import ( + model_args, +) pytestmark = [ pytest.mark.filterwarnings( @@ -892,6 +898,146 @@ 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_force_reshape_uses_dynamic_shapes_after_retracing() -> None: + """Generalized atom dimensions still normalize flattened force labels.""" + + @tf.function(reduce_retracing=True) + def normalize(force: tf.Tensor, label: tf.Tensor) -> tf.Tensor: + result = Trainer._match_label_shapes( + {"force": wrap_tensor(force)}, + {"force": label}, + ) + return to_tf_tensor(result["force"]) + + for natoms in (2, 3): + force = tf.reshape( + tf.range(3 * natoms, dtype=tf.float64), + (1, natoms, 3), + ) + normalized = normalize(force, tf.zeros((1, 3 * natoms), tf.float64)) + assert tuple(normalized.shape) == (1, 3 * natoms) + + +@pytest.mark.timeout(180) +def test_compiled_dpa4_training_step_accepts_two_atom_counts() -> None: + """Default degree-channel mixing is graph-safe under reduce_retracing.""" + model_params = model_args().normalize_value( + { + "type": "dpa4", + "type_map": ["A", "B"], + "descriptor": { + "type": "dpa4", + "sel": 4, + "rcut": 4.0, + "channels": 4, + "n_radial": 4, + "lmax": 1, + "mmax": 1, + "n_blocks": 1, + "radial_so2_mode": "degree_channel", + "random_gamma": False, + "precision": "float64", + "seed": 1, + }, + "fitting_net": { + "type": "dpa4_ener", + "neuron": [4], + "precision": "float64", + "seed": 1, + }, + }, + trim_pattern="_.*", + ) + model = get_model(model_params) + optimizer = tf.keras.optimizers.SGD(1.0e-3) + + @tf.function(reduce_retracing=True) + def train_step( + coord: tf.Tensor, + atype: tf.Tensor, + nlist: tf.Tensor, + mapping: tf.Tensor, + ) -> tuple[tf.Tensor, tf.Tensor]: + with tf.GradientTape() as tape: + result = model.call_common_lower( + coord, + atype, + nlist, + mapping, + None, + None, + ) + loss = tf.reduce_sum(to_tf_tensor(result["energy_redu"])) + gradients = tape.gradient(loss, model.trainable_variables) + optimizer.apply_gradients( + (gradient, variable) + for gradient, variable in zip( + gradients, model.trainable_variables, strict=True + ) + if gradient is not None + ) + return loss, to_tf_tensor(result["energy_derv_r"]) + + def dense_nlist(natoms: int) -> tf.Tensor: + rows = [] + for atom in range(natoms): + neighbors = [other for other in range(natoms) if other != atom] + rows.append(neighbors + [-1] * (4 - len(neighbors))) + return tf.constant([rows], dtype=tf.int64) + + for natoms in (2, 3): + coord = tf.reshape( + tf.range(3 * natoms, dtype=tf.float64) * 0.1, + (1, natoms, 3), + ) + _, force = train_step( + coord, + tf.zeros((1, natoms), dtype=tf.int32), + dense_nlist(natoms), + tf.range(natoms, dtype=tf.int64)[tf.newaxis, :], + ) + assert tuple(force.shape) == (1, natoms, 1, 3) + + def test_model_ret_translation_only_uses_label_virial_when_not_requested() -> None: trainer = object.__new__(Trainer) trainer.models = {