feat(tf2): support DPA4 training

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]

deepmd/tf2/common.py

@@ -2,6 +2,8 @@
 
 from collections.abc import (
     Callable,
+    Mapping,
+    Sequence,
 )
 from functools import (
     wraps,
@@ -104,6 +106,7 @@ def unwrap_value(value: Any) -> Any:
     f"{_PACKAGE_ROOT}.descriptor.dpa2",
     f"{_PACKAGE_ROOT}.descriptor.repflows",
     f"{_PACKAGE_ROOT}.descriptor.dpa3",
+    f"{_PACKAGE_ROOT}.descriptor.dpa4",
     f"{_PACKAGE_ROOT}.descriptor.hybrid",
     f"{_PACKAGE_ROOT}.fitting",
     f"{_PACKAGE_ROOT}.atomic_model.dp_atomic_model",
@@ -253,6 +256,47 @@ def tf2_module(module: type[T]) -> type[T]:
 
     @wraps(module, updated=())
     class TF2Module(module, tf.Module):  # type: ignore[misc, valid-type]
+        @staticmethod
+        def _tf2_array_variable_storage_name(name: str) -> str:
+            return f"_tf2_{name}_variable"
+
+        @staticmethod
+        def _tf2_array_variable_list_storage_name(name: str) -> str:
+            return f"_tf2_{name}_variables"
+
+        def _tf2_array_variable_attr_names(self) -> set[str]:
+            return set(getattr(self, "_tf2_array_variable_attrs", ()))
+
+        def _tf2_array_variable_list_attr_names(self) -> set[str]:
+            return set(getattr(self, "_tf2_array_variable_list_attrs", ()))
+
+        def _set_tf2_array_variable(self, name: str, value: Any) -> None:
+            storage_name = self._tf2_array_variable_storage_name(name)
+            if value is None:
+                tf.Module.__setattr__(self, storage_name, None)
+                return
+            tensor = to_tf_tensor(value)
+            variable = tf.Variable(
+                tensor,
+                trainable=bool(getattr(self, "trainable", True)),
+                name=name,
+            )
+            tf.Module.__setattr__(self, storage_name, variable)
+
+        def _set_tf2_array_variable_list(self, name: str, value: Any) -> None:
+            storage_name = self._tf2_array_variable_list_storage_name(name)
+            variables = []
+            for idx, item in enumerate(value):
+                tensor = to_tf_tensor(item)
+                variables.append(
+                    tf.Variable(
+                        tensor,
+                        trainable=bool(getattr(self, "trainable", True)),
+                        name=f"{name}_{idx}",
+                    )
+                )
+            tf.Module.__setattr__(self, storage_name, variables)
+
         def __init__(self, *args: Any, **kwargs: Any) -> None:
             tf.Module.__init__(self)
             super().__init__(*args, **kwargs)
@@ -266,11 +310,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):

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",