feat(tf2): support DPA4 descriptor

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

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