[Code scan] Handle multi-frame TensorFlow DeepSpin extension buffers #5660

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

Problem

The TensorFlow DeepSpin C++ API accepts multi-frame coordinate and spin inputs, but the spin virtual-atom extension and output copyback paths only handle the first frame.

The public API documents coord and spin as nframes x natoms x 3:

• deepmd-kit/source/api_cc/include/DeepSpin.h

  Lines 53 to 56 in 73de44b

  	* @param[in] coord The coordinates of atoms. The array should be of size
  	*nframes x natoms x 3.
  	* @param[in] spin The spins of atoms, [0, 0, 0] if no spin. The array should
  	*be of size nframes x natoms x 3.

The direct TensorFlow path infers nframes from the coordinate vector:

• deepmd-kit/source/api_cc/src/DeepSpinTF.cc

  Lines 593 to 596 in 73de44b

  	const bool atomic) {
  	// if datype.size is 0, not clear nframes; but 1 is just ok
  	int nframes = datype_.size() > 0 ? (dcoord_.size() / 3 / datype_.size()) : 1;
  	int nloc = datype_.size();

However, extend_nlist() builds only one frame of extended coordinates and indexes dcoord_/dspin_ without any frame offset:

• deepmd-kit/source/api_cc/src/DeepSpinTF.cc

  Lines 1222 to 1247 in 73de44b

  	// extend coord and atype
  	int nloc = datype_.size();
  	int nloc_spin = 0;
  	for (int ii = 0; ii < nloc; ii++) {
  	if (datype_[ii] < ntypes_spin) {
  	nloc_spin += 1;
  	}
  	}
  	int extend_nall = nloc + nloc_spin;
  	extend_dcoord.resize(static_cast<size_t>(extend_nall) * 3);
  	extend_atype.resize(extend_nall);
  	for (int ii = 0; ii < nloc; ii++) {
  	extend_atype[ii] = datype_[ii];
  	if (datype_[ii] < ntypes_spin) {
  	extend_atype[ii + nloc] = datype_[ii] + ntypes - ntypes_spin;
  	}
  	for (int jj = 0; jj < 3; jj++) {
  	extend_dcoord[ii * 3 + jj] = dcoord_[ii * 3 + jj];
  	if (datype_[ii] < ntypes_spin) {
  	extend_dcoord[(ii + nloc) * 3 + jj] =
  	dcoord_[ii * 3 + jj] + dspin_[ii * 3 + jj] /
  	spin_norm[datype_[ii]] *
  	virtual_len[datype_[ii]];
  	}
  	}
  	}

After inference, the direct path resizes force and magnetic-force outputs for all frames:

• deepmd-kit/source/api_cc/src/DeepSpinTF.cc

  Lines 636 to 638 in 73de44b

  	// backward force and mag.
  	dforce_.resize(static_cast<size_t>(nframes) * nloc * 3);
  	dforce_mag_.resize(static_cast<size_t>(nframes) * nloc * 3);

But it copies only frame-zero offsets:

• deepmd-kit/source/api_cc/src/DeepSpinTF.cc

  Lines 639 to 645 in 73de44b

  	for (int ii = 0; ii < nloc; ++ii) {
  	for (int dd = 0; dd < 3; ++dd) {
  	dforce_[3 * ii + dd] = dforce_tmp[3 * ii + dd];
  	if (datype_[ii] < ntypes_spin) {
  	dforce_mag_[3 * ii + dd] = dforce_tmp[3 * (ii + nloc) + dd];
  	} else {
  	dforce_mag_[3 * ii + dd] = 0.0;

The neighbor-list path has the same issue: it infers nframes, but extend() fills a single-frame extended coordinate buffer without frame offsets:

• deepmd-kit/source/api_cc/src/DeepSpinTF.cc

  Lines 729 to 738 in 73de44b

  	int nall = datype_.size();
  	// if nall==0, unclear nframes, but 1 is ok
  	int nframes = nall > 0 ? (dcoord_.size() / nall / 3) : 1;
  	int nloc = nall - nghost;
  	std::vector<VALUETYPE> extend_dcoord;
  	extend(extend_inum, extend_ilist, extend_numneigh, extend_neigh,
  	extend_firstneigh, extend_dcoord, extend_dtype, extend_nghost,
  	new_idx_map, old_idx_map, lmp_list, dcoord_, datype_, nghost, dspin_,
  	ntypes, ntypes_spin);

• deepmd-kit/source/api_cc/src/DeepSpinTF.cc

  Lines 1130 to 1150 in 73de44b

  	// extend coord
  	extend_dcoord.resize(static_cast<size_t>(extend_nall) * 3);
  	for (int ii = 0; ii < nloc; ii++) {
  	for (int jj = 0; jj < 3; jj++) {
  	extend_dcoord[new_idx_map[ii] * 3 + jj] = dcoord[ii * 3 + jj];
  	if (atype[ii] < numb_types_spin) {
  	double temp_dcoord = dcoord[ii * 3 + jj] + spin[ii * 3 + jj] /
  	spin_norm[atype[ii]] *
  	virtual_len[atype[ii]];
  	extend_dcoord[(new_idx_map[ii] + nloc) * 3 + jj] = temp_dcoord;
  	}
  	}
  	}
  	for (int ii = nloc; ii < nall; ii++) {
  	for (int jj = 0; jj < 3; jj++) {
  	extend_dcoord[new_idx_map[ii] * 3 + jj] = dcoord[ii * 3 + jj];
  	if (atype[ii] < numb_types_spin) {
  	double temp_dcoord = dcoord[ii * 3 + jj] + spin[ii * 3 + jj] /
  	spin_norm[atype[ii]] *
  	virtual_len[atype[ii]];
  	extend_dcoord[(new_idx_map[ii] + nghost) * 3 + jj] = temp_dcoord;

Its copyback also sizes outputs for all frames but writes only frame-zero offsets for force, magnetic force, atom energy, and atom virial:

• deepmd-kit/source/api_cc/src/DeepSpinTF.cc

  Lines 807 to 827 in 73de44b

  	dforce_.resize(static_cast<size_t>(nframes) * nall * 3);
  	dforce_mag_.resize(static_cast<size_t>(nframes) * nall * 3);
  	datom_energy_.resize(static_cast<size_t>(nframes) * nall);
  	datom_virial_.resize(static_cast<size_t>(nframes) * nall * 9);
  	for (int ii = 0; ii < nall; ++ii) {
  	int new_idx = new_idx_map[ii];
  	for (int dd = 0; dd < 3; ++dd) {
  	dforce_[3 * ii + dd] = dforce_tmp[3 * new_idx + dd];
  	datom_energy_[ii] = datom_energy_tmp[new_idx];
  	if (datype_[ii] < ntypes_spin && ii < nloc) {
  	dforce_mag_[3 * ii + dd] = dforce_tmp[3 * (new_idx + nloc) + dd];
  	} else if (datype_[ii] < ntypes_spin) {
  	dforce_mag_[3 * ii + dd] = dforce_tmp[3 * (new_idx + nghost) + dd];
  	} else {
  	dforce_mag_[3 * ii + dd] = 0.0;
  	}
  	}
  	for (int dd = 0; dd < 9; ++dd) {
  	datom_virial_[ii * 9 + dd] = datom_virial_tmp[new_idx * 9 + dd];
  	}

Impact

Multi-frame DeepSpinTF calls can ignore later-frame coordinates/spins during virtual-atom construction and leave later-frame outputs uninitialized or incorrectly copied. This can produce wrong forces, magnetic forces, atomic energies, and atomic virials for nframes > 1.

Suggested fix

Make extend() and extend_nlist() build extended coordinate/type buffers for every frame, and copy outputs back with explicit frame * natoms * stride offsets. Add direct and neighbor-list TensorFlow DeepSpin regressions with nframes > 1 and distinct spins in each frame.