From 86b61b18090511ca0cd8c89388836aea3e1054b6 Mon Sep 17 00:00:00 2001 From: njzjz-bot Date: Fri, 17 Jul 2026 05:43:03 +0800 Subject: [PATCH 1/3] 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 --- source/op/tf/prod_force_grad.cc | 25 +++++++--- source/tests/tf/test_prod_force_grad.py | 64 +++++++++++++++++++++++++ 2 files changed, 83 insertions(+), 6 deletions(-) diff --git a/source/op/tf/prod_force_grad.cc b/source/op/tf/prod_force_grad.cc index f5b9295643..b5704038cd 100644 --- a/source/op/tf/prod_force_grad.cc +++ b/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(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(nloc) * ndescrpt * 12 == in_deriv_shape.dim_size(1)), @@ -118,10 +124,20 @@ class ProdForceGradOp : public OpKernel { auto axis = axis_tensor.flat(); auto grad_net = grad_net_tensor->flat(); + // 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(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; } diff --git a/source/tests/tf/test_prod_force_grad.py b/source/tests/tf/test_prod_force_grad.py index 6fea4a631d..373e870f18 100644 --- a/source/tests/tf/test_prod_force_grad.py +++ b/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__() From 9a66c055b1123e9409c463be20c31952a06dd670 Mon Sep 17 00:00:00 2001 From: njzjz-bot Date: Mon, 27 Jul 2026 12:06:33 +0800 Subject: [PATCH 2/3] perf(tf): bound ProdForceGrad neighbors with one reduction The per-element OP_REQUIRES pre-pass ran serially outside the omp region on every backward step. Replace it with a parallel reduction and a single check, and note that the shape guards document the training-only contract rather than a reachable failure. --- source/op/tf/prod_force_grad.cc | 18 +++++++++++++++--- 1 file changed, 15 insertions(+), 3 deletions(-) diff --git a/source/op/tf/prod_force_grad.cc b/source/op/tf/prod_force_grad.cc index b5704038cd..18ee11cfd3 100644 --- a/source/op/tf/prod_force_grad.cc +++ b/source/op/tf/prod_force_grad.cc @@ -86,6 +86,10 @@ class ProdForceGradOp : public OpKernel { context, (nframes == axis_shape.dim_size(0)), deepmd::tf_compat::InvalidArgument("number of frames should match")); + // natoms is [nloc, nall, ...]. In training nall == nloc, so these two + // checks are contract documentation rather than a reachable failure: the + // only caller is TF's registered gradient for ProdForce, which never sees + // ghosts. They pin the layout grad() is indexed with further down. OP_REQUIRES( context, (nall >= nloc), deepmd::tf_compat::InvalidArgument( @@ -127,12 +131,20 @@ class ProdForceGradOp : public OpKernel { // 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 ever sees the neighbor list that ProdForce + // consumed, so an index past nall cannot arise there. This op is public, + // though, and j_idx addresses grad() directly below, so bound it once with + // a parallel reduction rather than a per-element check on the hot path. const int64_t nlist_size = static_cast(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) { - OP_REQUIRES(context, nlist(ii) < nall, - deepmd::tf_compat::InvalidArgument( - "neighbor index should be smaller than all atoms")); + 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 From 3e4d64fd9d49bbaa49e3a58b960afacca1b76b93 Mon Sep 17 00:00:00 2001 From: njzjz-bot Date: Thu, 30 Jul 2026 11:35:17 +0800 Subject: [PATCH 3/3] docs(tf): clarify ProdForceGrad raw-op contract Address the outstanding requested-change review comments. Coding-Agent: Codex Codex-Version: codex-cli 0.144.6 Model: gpt-5.6-sol Reasoning-Effort: xhigh --- source/op/tf/prod_force_grad.cc | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/source/op/tf/prod_force_grad.cc b/source/op/tf/prod_force_grad.cc index 18ee11cfd3..2cdc010701 100644 --- a/source/op/tf/prod_force_grad.cc +++ b/source/op/tf/prod_force_grad.cc @@ -86,10 +86,10 @@ class ProdForceGradOp : public OpKernel { context, (nframes == axis_shape.dim_size(0)), deepmd::tf_compat::InvalidArgument("number of frames should match")); - // natoms is [nloc, nall, ...]. In training nall == nloc, so these two - // checks are contract documentation rather than a reachable failure: the - // only caller is TF's registered gradient for ProdForce, which never sees - // ghosts. They pin the layout grad() is indexed with further down. + // 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( @@ -132,10 +132,10 @@ class ProdForceGradOp : public OpKernel { // those upstream gradients distinct: folding ghost indices modulo nloc // would differentiate a different output than the forward op produced. // - // The registered gradient only ever sees the neighbor list that ProdForce - // consumed, so an index past nall cannot arise there. This op is public, - // though, and j_idx addresses grad() directly below, so bound it once with - // a parallel reduction rather than a per-element check on the hot path. + // 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(nframes) * nloc * nnei; int nlist_out_of_range = 0; #pragma omp parallel for reduction(| : nlist_out_of_range)