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

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)
             )

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

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

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,11 +307,17 @@ 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
+        # MLPLayer stores trainability as serialization metadata but its
+        # Parameters default to requires_grad=True. Apply the per-layer policy
+        # to the actual tensors so construction and deserialization agree.
+        for layer_idx, layer in enumerate(self.hidden_layers):
+            for parameter in layer.parameters():
+                parameter.requires_grad = self.trainable[layer_idx]
+        for parameter in self.output_layer.parameters():
+            parameter.requires_grad = self.trainable[-1]
 
     def _apply_input_film(
         self,
@@ -395,6 +408,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()},
         }