This issue comes from a Codex global scan of deepmodeling/deepmd-kit at commit 73de44b1f94471b2e3bdb6b11f57b34d7bc791bb.

Problem

The legacy TensorFlow ProdForceGrad path handles the first ghost atom inconsistently and can read past the local force-gradient buffer.

ProdForce outputs forces for nall atoms:

• deepmd-kit/source/op/tf/prod_force.cc

  Lines 90 to 98 in 73de44b

  	// Create an output tensor
  	TensorShape force_shape;
  	force_shape.AddDim(nframes);
  	force_shape.AddDim(3 * static_cast<int64_t>(nall));
  	// std::cout << "forcesahpe " << force_shape.dim_size(0) << " " <<
  	// force_shape.dim_size(1) << std::endl;
  	Tensor* force_tensor = NULL;
  	OP_REQUIRES_OK(context,
  	context->allocate_output(0, force_shape, &force_tensor));

Its registered Python gradient passes TensorFlow's grad tensor directly into prod_force_grad:

• deepmd-kit/deepmd/tf/op/_prod_force_grad.py

  Lines 15 to 26 in 73de44b

  	@ops.RegisterGradient("ProdForce")
  	def _prod_force_grad_cc(op: tf.Operation, grad: tf.Tensor) -> list[tf.Tensor | None]:
  	net_grad = op_grads_module.prod_force_grad(
  	grad,
  	op.inputs[0],
  	op.inputs[1],
  	op.inputs[2],
  	op.inputs[3],
  	op.inputs[4],
  	n_a_sel=op.get_attr("n_a_sel"),
  	n_r_sel=op.get_attr("n_r_sel"),
  	)

But the C++ grad op requires grad to have only 3 * nloc entries:

• deepmd-kit/source/op/tf/prod_force_grad.cc

  Lines 74 to 90 in 73de44b

  	// check the sizes
  	OP_REQUIRES(
  	context, (nframes == grad_shape.dim_size(0)),
  	deepmd::tf_compat::InvalidArgument("number of frames should match"));
  	OP_REQUIRES(
  	context, (nframes == in_deriv_shape.dim_size(0)),
  	deepmd::tf_compat::InvalidArgument("number of frames should match"));
  	OP_REQUIRES(
  	context, (nframes == nlist_shape.dim_size(0)),
  	deepmd::tf_compat::InvalidArgument("number of frames should match"));
  	OP_REQUIRES(
  	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)),
  	deepmd::tf_compat::InvalidArgument(
  	"input grad shape should be 3 x natoms"));

When processing neighbor indices, it remaps only j_idx > nloc:

• deepmd-kit/source/op/tf/prod_force_grad.cc

  Lines 163 to 168 in 73de44b

  	// 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;
  	}

The first ghost atom has index j_idx == nloc, so it is not remapped and later reads grad + nloc * 3, one element past the accepted local gradient buffer:

• deepmd-kit/source/op/tf/prod_force_grad.cc

  Lines 172 to 198 in 73de44b

  	if (jj == axis_0) {
  	for (int aa = 0; aa < ndescrpt; ++aa) {
  	for (int dd = 0; dd < 3; ++dd) {
  	grad_net(grad_net_iter + i_idx * ndescrpt + aa) -=
  	grad(grad_iter + j_idx * 3 + dd) *
  	in_deriv(in_iter + i_idx * ndescrpt * 12 + aa * 12 + 3 +
  	dd);
  	}
  	}
  	} else if (jj == axis_1) {
  	for (int aa = 0; aa < ndescrpt; ++aa) {
  	for (int dd = 0; dd < 3; ++dd) {
  	grad_net(grad_net_iter + i_idx * ndescrpt + aa) -=
  	grad(grad_iter + j_idx * 3 + dd) *
  	in_deriv(in_iter + i_idx * ndescrpt * 12 + aa * 12 + 6 +
  	dd);
  	}
  	}
  	} else {
  	int aa_start, aa_end;
  	make_descript_range(aa_start, aa_end, jj);
  	for (int aa = aa_start; aa < aa_end; ++aa) {
  	for (int dd = 0; dd < 3; ++dd) {
  	grad_net(grad_net_iter + i_idx * ndescrpt + aa) -=
  	grad(grad_iter + j_idx * 3 + dd) *
  	in_deriv(in_iter + i_idx * ndescrpt * 12 + aa * 12 + 9 +
  	dd);

The shared force-gradient helper uses >= nloc for the same boundary:

• deepmd-kit/source/lib/src/prod_force_grad.cc

  Lines 53 to 69 in 73de44b

  	// loop over neighbors
  	for (int jj = 0; jj < nnei; ++jj) {
  	int j_idx = nlist[i_idx * nnei + jj];
  	if (j_idx >= nloc) {
  	j_idx = j_idx % nloc;
  	}
  	if (j_idx < 0) {
  	continue;
  	}
  	int aa_start, aa_end;
  	make_index_range(aa_start, aa_end, jj, nnei);
  	const int kk = i_idx / nloc; // frame index
  	for (int aa = aa_start; aa < aa_end; ++aa) {
  	for (int dd = 0; dd < 3; ++dd) {
  	grad_net[i_idx * ndescrpt + aa] +=
  	grad[kk * nloc * 3 + j_idx * 3 + dd] *
  	env_deriv[i_idx * ndescrpt * 3 + aa * 3 + dd];

Impact

For loc-frame models with nall > nloc, the legacy ProdForce gradient can reject TensorFlow's natural 3 * nall gradient shape or read past a local-only gradient when the neighbor list contains the first ghost atom.

Suggested fix

Align the gradient op contract with the forward output shape, or if the intended contract is local-only, remap ghost indices with j_idx >= nloc consistently. Add a loc-frame gradient regression with nall > nloc and a neighbor at exactly nloc.