fix(tf): preserve ghost gradients in ProdForceGrad

Size upstream force gradients by all atoms and keep local and ghost slices distinct in the legacy TensorFlow ProdForce gradient. Validate atom counts and neighbor bounds before indexing.

Cover a second ghost at index nloc + 1 so the regression exercises the removed modulo branch and proves its independent upstream gradient is used.

Coding-Agent: Codex
Codex-Version: codex-cli 0.144.4
Model: gpt-5.6-sol
Reasoning-Effort: xhigh

Author: njzjz-bot

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,14 @@ 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)),
+    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 +124,20 @@ 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.
+    const int64_t nlist_size = static_cast<int64_t>(nframes) * nloc * nnei;
+    for (int64_t ii = 0; ii < nlist_size; ++ii) {
+      OP_REQUIRES(context, nlist(ii) < nall,
+                  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 +179,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__()