fix(tf): align ProdForceGrad with ProdForce layout

source/op/tf/prod_force_grad.cc

@@ -68,6 +68,7 @@ class ProdForceGradOp : public OpKernel {
 
     int nframes = net_deriv_tensor.shape().dim_size(0);
     int nloc = natoms(0);
+    int nall = natoms(1);
     int ndescrpt = nloc > 0 ? net_deriv_tensor.shape().dim_size(1) / nloc : 0;
     int nnei = nloc > 0 ? nlist_tensor.shape().dim_size(1) / nloc : 0;
 
@@ -85,9 +86,18 @@ class ProdForceGradOp : public OpKernel {
         context, (nframes == axis_shape.dim_size(0)),
         deepmd::tf_compat::InvalidArgument("number of frames should match"));
 
-    OP_REQUIRES(context, (nloc * 3 == grad_shape.dim_size(1)),
+    // natoms is [nloc, nall, ...]. The registered training gradient normally
+    // has nall == nloc, while callers of the public raw op can supply an
+    // extended layout. These checks enforce the layout grad() is indexed with
+    // for both entry points.
+    OP_REQUIRES(
+        context, (nall >= nloc),
+        deepmd::tf_compat::InvalidArgument(
+            "number of all atoms should not be smaller than local atoms"));
+    OP_REQUIRES(context,
+                (static_cast<int64_t>(nall) * 3 == grad_shape.dim_size(1)),
                 deepmd::tf_compat::InvalidArgument(
-                    "input grad shape should be 3 x natoms"));
+                    "input grad shape should be 3 x all atoms"));
     OP_REQUIRES(context,
                 (static_cast<int64_t>(nloc) * ndescrpt * 12 ==
                  in_deriv_shape.dim_size(1)),
@@ -118,10 +128,28 @@ class ProdForceGradOp : public OpKernel {
     auto axis = axis_tensor.flat<int>();
     auto grad_net = grad_net_tensor->flat<FPTYPE>();
 
+    // ProdForce returns one force vector for every local and ghost atom. Keep
+    // those upstream gradients distinct: folding ghost indices modulo nloc
+    // would differentiate a different output than the forward op produced.
+    //
+    // The registered gradient only sees the neighbor list that ProdForce
+    // consumed, but the raw op is public and j_idx addresses grad() directly
+    // below. Bound it once with a parallel reduction rather than a per-element
+    // check on the hot path.
+    const int64_t nlist_size = static_cast<int64_t>(nframes) * nloc * nnei;
+    int nlist_out_of_range = 0;
+#pragma omp parallel for reduction(| : nlist_out_of_range)
+    for (int64_t ii = 0; ii < nlist_size; ++ii) {
+      nlist_out_of_range |= (nlist(ii) >= nall);
+    }
+    OP_REQUIRES(context, (!nlist_out_of_range),
+                deepmd::tf_compat::InvalidArgument(
+                    "neighbor index should be smaller than all atoms"));
+
     // loop over frames
 #pragma omp parallel for
     for (int kk = 0; kk < nframes; ++kk) {
-      int grad_iter = kk * nloc * 3;
+      int grad_iter = kk * nall * 3;
       int net_iter = kk * nloc * ndescrpt;
       int in_iter = kk * nloc * ndescrpt * 12;
       int nlist_iter = kk * nloc * nnei;
@@ -163,9 +191,6 @@ class ProdForceGradOp : public OpKernel {
         // loop over neighbors
         for (int jj = 0; jj < nnei; ++jj) {
           int j_idx = nlist(nlist_iter + i_idx * nnei + jj);
-          if (j_idx > nloc) {
-            j_idx = j_idx % nloc;
-          }
           if (j_idx < 0) {
             continue;
           }

source/tests/tf/test_prod_force_grad.py

@@ -4,10 +4,74 @@
 from deepmd.tf.env import (
     GLOBAL_TF_FLOAT_PRECISION,
     op_grads_module,
+    op_module,
     tf,
 )
 
 
+class TestLegacyProdForceGradGhost(tf.test.TestCase):
+    """Exercise a non-boundary ghost through ProdForce's registered gradient."""
+
+    def test_second_ghost_uses_its_own_upstream_gradient(self) -> None:
+        nloc = 1
+        nall = 3
+        nframes = 2
+        ndescrpt = 4
+        net_deriv = tf.placeholder(tf.float64, [nframes, nloc * ndescrpt])
+
+        in_deriv = np.zeros((nframes, nloc * ndescrpt * 12))
+        # The only neighbor is axis 0. Give descriptor 0 a simple derivative
+        # with respect to the ghost atom and leave every other term zero.
+        in_deriv[:, 3:6] = [1.0, 2.0, 3.0]
+        force = op_module.prod_force(
+            net_deriv,
+            tf.constant(in_deriv, dtype=tf.float64),
+            # The second ghost has index nloc + 1. This specifically exercises
+            # the removed ``j_idx > nloc`` modulo branch; the first ghost at
+            # exactly nloc would never have entered that branch.
+            tf.constant([[nloc + 1], [nloc + 1]], dtype=tf.int32),
+            tf.constant([[0, 0, 0, 0], [0, 0, 0, 0]], dtype=tf.int32),
+            tf.constant([nloc, nall, nloc], dtype=tf.int32),
+            n_a_sel=1,
+            n_r_sel=0,
+        )
+
+        # Local and both ghost gradients deliberately differ. The expected
+        # result must use the second-ghost slice, not fold index 2 onto local
+        # index 0 as the old modulo branch did.
+        upstream = tf.constant(
+            [
+                [10.0, 11.0, 12.0, 7.0, 8.0, 9.0, 4.0, 5.0, 6.0],
+                [20.0, 21.0, 22.0, 8.0, 7.0, 6.0, 1.0, 2.0, 3.0],
+            ],
+            dtype=tf.float64,
+        )
+        grad_net = tf.gradients(force, net_deriv, grad_ys=upstream)[0]
+
+        with self.cached_session() as sess:
+            actual_force, actual_grad = sess.run(
+                [force, grad_net],
+                feed_dict={
+                    net_deriv: [
+                        [1.0, 0.0, 0.0, 0.0],
+                        [2.0, 0.0, 0.0, 0.0],
+                    ]
+                },
+            )
+
+        np.testing.assert_array_equal(
+            actual_force,
+            [
+                [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, -1.0, -2.0, -3.0],
+                [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, -2.0, -4.0, -6.0],
+            ],
+        )
+        np.testing.assert_array_equal(
+            actual_grad,
+            [[-32.0, 0.0, 0.0, 0.0], [-14.0, 0.0, 0.0, 0.0]],
+        )
+
+
 class TestProdForceGrad(tf.test.TestCase):
     def setUp(self) -> None:
         self.sess = self.cached_session().__enter__()