From 86b61b18090511ca0cd8c89388836aea3e1054b6 Mon Sep 17 00:00:00 2001 From: njzjz-bot Date: Fri, 17 Jul 2026 05:43:03 +0800 Subject: [PATCH] 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__()