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__()