diff --git a/deepmd/dpmodel/array_api.py b/deepmd/dpmodel/array_api.py index ca597b3444..0db1c51603 100644 --- a/deepmd/dpmodel/array_api.py +++ b/deepmd/dpmodel/array_api.py @@ -201,6 +201,22 @@ 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,)) + values_tensor = values.unwrap() + # unsorted_segment_sum rather than scatter_nd: both accumulate repeated + # indices, but scatter_nd rejects a destination with no elements even + # when the updates are empty too, which a descriptor call with zero + # edges legitimately produces. + updates = tf.math.unsorted_segment_sum( + values_tensor, + indices_tensor, + tf.shape(x_tensor, out_type=tf.int64)[0], + ) + 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 +286,52 @@ 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], + ) + if values_tensor.dtype.is_floating: + # TensorFlow uses the lowest finite value as the identity of + # unsorted_segment_max. Restore the true maximum-at identity when + # every update for a touched segment element is negative infinity. + all_negative_infinity = ( + tf.math.unsorted_segment_min( + tf.cast( + tf.math.is_inf(values_tensor) & (values_tensor < 0), + tf.int32, + ), + indices_tensor, + tf.shape(x_tensor, out_type=tf.int64)[0], + ) + > 0 + ) + reduced = tf.where( + all_negative_infinity, + tf.cast(float("-inf"), values_tensor.dtype), + 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] @@ -337,12 +399,12 @@ def xp_setitem_at(x: Array, mask: Array, values: Array) -> Array: def xp_uniform(like: Array, size: int, low: float = 0.0, high: float = 1.0) -> Array: """Draw ``size`` uniform samples in ``[low, high)`` on ``like``'s device. - Each backend uses its own generator: torch draws with ``torch.rand`` (as - pt does, so ``setup_seed`` replays it, with no host copy -- and a host - draw would freeze to a constant under tracing); other backends use - :mod:`deepmd.utils.random`, which ``setup_seed`` also seeds. Draws are - therefore not comparable across backends -- use only for a per-forward - random stream, never where a parity test looks. + Each backend uses its own generator: TensorFlow draws with + ``tf.random.uniform`` so traced graphs advance the runtime RNG, torch draws + with ``torch.rand`` (so ``setup_seed`` replays it without a host copy), and + other backends use :mod:`deepmd.utils.random`. Draws are therefore not + comparable across backends -- use only for a per-forward random stream, + never where a parity test looks. Parameters ---------- @@ -360,6 +422,18 @@ def xp_uniform(like: Array, size: int, low: float = 0.0, high: float = 1.0) -> A Array Samples of shape ``(size,)`` matching ``like``. """ + xp = array_api_compat.array_namespace(like) + if getattr(xp, "__name__", "") == "deepmd._vendors.ndtensorflow": + import tensorflow as tf + + sample_shape = tf.reshape(tf.cast(size, tf.int32), (1,)) + samples = tf.random.uniform( + sample_shape, + minval=low, + maxval=high, + dtype=like.dtype, + ) + return xp.asarray(samples) if array_api_compat.is_torch_array(like): import torch @@ -368,7 +442,6 @@ def xp_uniform(like: Array, size: int, low: float = 0.0, high: float = 1.0) -> A ) from deepmd.utils import random as dp_random - xp = array_api_compat.array_namespace(like) drawn = np.asarray(dp_random.random(size)) * (high - low) + low return xp.astype( xp_asarray_nodetach(xp, drawn, device=array_api_compat.device(like)), diff --git a/deepmd/dpmodel/descriptor/dpa4_nn/so2.py b/deepmd/dpmodel/descriptor/dpa4_nn/so2.py index 4c8cd80e9a..83bb38297a 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/tf2/common.py b/deepmd/tf2/common.py index bb8155a38c..dba06f3874 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,72 @@ 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", {} + ) + if value is None: + tf.Module.__setattr__(self, storage_name, None) + object.__getattribute__(self, "__dict__").pop(name, None) + return + 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 +335,117 @@ 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 ( + None + if variables is None + else [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 None + if name in self._tf2_array_variable_list_attr_names(): + self._set_tf2_array_variable_list(name, value) + return None 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..d41c188867 --- /dev/null +++ b/deepmd/tf2/descriptor/dpa4.py @@ -0,0 +1,325 @@ +# 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) + self._tf2_training_mode = False + _promote_trainable_tree(self) + + @classmethod + def deserialize(cls, data: dict) -> "DescrptDPA4": + obj = super().deserialize(data) + return _promote_trainable_tree(obj) + + def set_training_mode(self, training: bool) -> None: + """Select train/eval graph tracing for training-only augmentation.""" + self._tf2_training_mode = bool(training) + + def _in_training_mode(self) -> bool: + """Return the Python trace-time training state used by TF2 graphs.""" + return self._tf2_training_mode 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/train/trainer.py b/deepmd/tf2/train/trainer.py index 6fde31e8f8..4c1e8af671 100644 --- a/deepmd/tf2/train/trainer.py +++ b/deepmd/tf2/train/trainer.py @@ -959,6 +959,7 @@ def compiled_prepared_train_step( extended_coord_corr, label_dict=label_dict, do_virial=do_virial, + training=True, ) loss, more_loss = self.losses[task_key]( learning_rate=cur_lr, @@ -1025,6 +1026,7 @@ def compiled_train_step( input_dict, label_dict=label_dict, do_virial=do_virial, + training=True, ) loss, more_loss = self.losses[task_key]( learning_rate=cur_lr, @@ -1154,6 +1156,7 @@ def compiled_eval_step( input_dict, label_dict=label_dict, do_virial=do_virial, + training=False, ) _, more_loss = self.losses[task_key]( learning_rate=cur_lr, @@ -1323,8 +1326,10 @@ def _call_model( *, label_dict: dict[str, Any] | None = None, do_virial: bool = True, + training: bool, ) -> dict[str, Any]: model = self.models[task_key] + self._set_model_training_mode(model, training) call_common = getattr(model, "call_common", None) if callable(call_common): model_ret = call_common( @@ -1366,8 +1371,10 @@ def _call_prepared_model( *, label_dict: dict[str, Any] | None = None, do_virial: bool = True, + training: bool, ) -> dict[str, Any]: model = self.models[task_key] + self._set_model_training_mode(model, training) call_lower_formatted = getattr(model, "_call_common_lower_formatted", None) if callable(call_lower_formatted): model_ret_lower = wrap_value( @@ -1413,6 +1420,22 @@ def _call_prepared_model( do_virial=do_virial, ) + @staticmethod + def _set_model_training_mode(model: Any, training: bool) -> None: + """Set descriptor train/eval state before TensorFlow traces a graph. + + TF2 models are ``tf.Module`` objects rather than Keras layers, so there + is no implicit ``training`` argument. The compiled train and validation + functions are traced separately; setting this Python flag at their call + seam lets DPA4 include graph-safe random-gamma augmentation only in the + training graph while keeping validation and exported inference stable. + """ + atomic_model = getattr(model, "atomic_model", None) + descriptor = getattr(atomic_model, "descriptor", None) + setter = getattr(descriptor, "set_training_mode", None) + if callable(setter): + setter(training) + def _translate_model_ret_to_loss_dict( self, task_key: str, diff --git a/source/tests/consistent/descriptor/test_dpa4.py b/source/tests/consistent/descriptor/test_dpa4.py index d059d5ffdb..38666106ec 100644 --- a/source/tests/consistent/descriptor/test_dpa4.py +++ b/source/tests/consistent/descriptor/test_dpa4.py @@ -23,6 +23,7 @@ INSTALLED_JAX, INSTALLED_PT, INSTALLED_PT_EXPT, + INSTALLED_TF2, CommonTest, parameterized_cases, ) @@ -38,6 +39,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_JAX: from deepmd.jax.descriptor.dpa4 import DescrptDPA4 as DescrptDPA4JAX else: @@ -164,12 +169,14 @@ def skip_pt(self) -> bool: skip_dp = False skip_tf = True + skip_tf2 = not INSTALLED_TF2 or DescrptDPA4TF2 is None skip_jax = not INSTALLED_JAX or DescrptDPA4JAX is None 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 @@ -248,6 +255,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_jax(self, jax_obj: Any) -> Any: return self.eval_jax_descriptor( jax_obj, diff --git a/source/tests/consistent/test_array_api.py b/source/tests/consistent/test_array_api.py index 05d88f3cc6..a9f2d8eaba 100644 --- a/source/tests/consistent/test_array_api.py +++ b/source/tests/consistent/test_array_api.py @@ -7,6 +7,7 @@ from deepmd.dpmodel.array_api import ( xp_add_at, xp_bincount, + xp_maximum_at, xp_scatter_sum, xp_setitem_at, xp_sigmoid, @@ -56,6 +57,21 @@ def test_torch_parameter_requires_grad(self) -> None: self.assertEqual(param.device, DEVICE) +class TestXpMaximumAtConsistent(unittest.TestCase): + """Test maximum-at identities that differ between backend primitives.""" + + @unittest.skipUnless(INSTALLED_TF2, "TensorFlow is not installed") + def test_tf_preserves_all_negative_infinity_segment(self) -> None: + x = tnp.asarray(np.full(2, -np.inf, dtype=np.float64)) + indices = tnp.asarray(np.array([0, 0], dtype=np.int64)) + values = tnp.asarray(np.array([-np.inf, -np.inf], dtype=np.float64)) + + result = to_numpy_array(xp_maximum_at(x, indices, values)) + + self.assertTrue(np.isneginf(result[0])) + self.assertTrue(np.isneginf(result[1])) + + class TestXpScatterSumConsistent(unittest.TestCase): """Test xp_scatter_sum consistency across backends.""" diff --git a/source/tests/tf2/test_dpa4.py b/source/tests/tf2/test_dpa4.py new file mode 100644 index 0000000000..8dd1670449 --- /dev/null +++ b/source/tests/tf2/test_dpa4.py @@ -0,0 +1,309 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Focused tests for TF2 DPA4 descriptor 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._vendors import ndtensorflow as xp +from deepmd.dpmodel.array_api import ( + xp_uniform, +) +from deepmd.dpmodel.utils.neighbor_graph import ( + NeighborGraph, +) +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, +) + + +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 + + +def test_promoted_optional_parameter_lists_accept_none() -> None: + """Disabled optional parameter lists retain their None sentinel.""" + descriptor = _make_trainable_descriptor() + promoted_lists = [] + for module in _iter_object_tree(descriptor): + for name in getattr(module, "_tf2_array_variable_list_attrs", ()): + promoted_lists.append((module, name)) + + assert promoted_lists + for module, name in promoted_lists: + setattr(module, name, None) + assert getattr(module, name) is None + storage_name = module._tf2_array_variable_list_storage_name(name) + assert object.__getattribute__(module, storage_name) is None + assert name not in object.__getattribute__(module, "__dict__") + + +def test_random_gamma_supports_train_and_eval_tracing_modes() -> None: + """Default DPA4 configs keep augmentation train-only in TF2.""" + descriptor = 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, + ) + + assert descriptor._in_training_mode() is False + descriptor.set_training_mode(True) + assert descriptor._in_training_mode() is True + restored = DescrptDPA4.deserialize(descriptor.serialize()) + assert restored.random_gamma is True + assert restored._in_training_mode() is False + + +def test_random_gamma_rng_advances_inside_tf_function() -> None: + """TensorFlow random draws must happen at execution, not trace, time.""" + like = xp.ones((16,), dtype=xp.float64) + + @tf.function + def draw_gamma(): + return xp_uniform(like, 16, 0.0, 2.0 * np.pi).unwrap() + + first = draw_gamma().numpy() + second = draw_gamma().numpy() + + assert not np.array_equal(first, second) + + +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) + + +def _make_edge_descriptor() -> DescrptDPA4: + """Build the smallest descriptor that still runs one interaction block.""" + 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=True, + seed=20260725, + ) + + +def test_call_graph_supports_an_empty_edge_list() -> None: + """A graph whose atom has no neighbors must reduce without scattering.""" + descriptor = _make_edge_descriptor() + + def run(n_node, atype, edge_index, edge_vec, edge_mask): + graph = NeighborGraph( + n_node=wrap_tensor(n_node), + edge_index=wrap_tensor(edge_index), + edge_vec=wrap_tensor(edge_vec), + edge_mask=wrap_tensor(edge_mask), + ) + descrpt, _ = descriptor.call_graph(graph, wrap_tensor(atype)) + return to_tf_tensor(descrpt) + + args = ( + to_tf_tensor(np.ones((1,), dtype=np.int64)), + to_tf_tensor(np.zeros((1,), dtype=np.int64)), + to_tf_tensor(np.zeros((2, 0), dtype=np.int64)), + to_tf_tensor(np.zeros((0, 3), dtype=np.float64)), + to_tf_tensor(np.zeros((0,), dtype=bool)), + ) + eager_descrpt = run(*args) + assert tuple(eager_descrpt.shape) == (1, descriptor.channels) + assert np.all(np.isfinite(eager_descrpt.numpy())) + + traced_descrpt = tf.function(run, reduce_retracing=True)(*args) + np.testing.assert_allclose(traced_descrpt.numpy(), eager_descrpt.numpy()) + + +def test_call_supports_a_zero_width_neighbor_list() -> None: + """The dense path must survive a neighbor list with no neighbor slots.""" + descriptor = _make_edge_descriptor() + coord_ext = to_tf_tensor(np.zeros((1, 1, 3), dtype=np.float64)) + atype_ext = to_tf_tensor(np.zeros((1, 1), dtype=np.int64)) + nlist = to_tf_tensor(np.zeros((1, 1, 0), dtype=np.int64)) + + def run(coord, atype, neighbors): + return to_tf_tensor( + descriptor(wrap_tensor(coord), wrap_tensor(atype), wrap_tensor(neighbors))[ + 0 + ] + ) + + eager = run(coord_ext, atype_ext, nlist) + assert tuple(eager.shape) == (1, 1, descriptor.channels) + assert np.all(np.isfinite(eager.numpy())) + + traced = tf.function(run, reduce_retracing=True)(coord_ext, atype_ext, nlist) + np.testing.assert_allclose(traced.numpy(), eager.numpy()) diff --git a/source/tests/tf2/test_training.py b/source/tests/tf2/test_training.py index 0c0a15b2b2..ee411b5373 100644 --- a/source/tests/tf2/test_training.py +++ b/source/tests/tf2/test_training.py @@ -633,6 +633,7 @@ def translated_output_def(self) -> dict[str, Any]: }, label_dict={"virial": tf.zeros((1, 9), dtype=tf.float64)}, do_virial=True, + training=False, ) assert captured == {