feat(c++): add DPA3 .pt2 export and C/C++ inference tests

Fix torch.export-incompatible slicing in dpmodel DPA3/repflows code
([:, :nloc] → xp_take_first_n, reshape with -1 → expand_dims), add
gen_dpa3.py model generation script, and create C++ test files for
both .pth and .pt2 backends with PBC and NoPbc fixtures.

Author: Han Wang

deepmd/dpmodel/descriptor/dpa3.py

@@ -10,6 +10,7 @@
 )
 from deepmd.dpmodel.array_api import (
     Array,
+    xp_take_first_n,
 )
 from deepmd.dpmodel.common import (
     cast_precision,
@@ -617,15 +618,19 @@ def call(
         type_embedding = self.type_embedding.call()
         if self.use_loc_mapping:
             node_ebd_ext = xp.reshape(
-                xp.take(type_embedding, xp.reshape(atype_ext[:, :nloc], (-1,)), axis=0),
+                xp.take(
+                    type_embedding,
+                    xp.reshape(xp_take_first_n(atype_ext, 1, nloc), (-1,)),
+                    axis=0,
+                ),
                 (nframes, nloc, self.tebd_dim),
             )
         else:
             node_ebd_ext = xp.reshape(
                 xp.take(type_embedding, xp.reshape(atype_ext, (-1,)), axis=0),
                 (nframes, nall, self.tebd_dim),
             )
-        node_ebd_inp = node_ebd_ext[:, :nloc, :]
+        node_ebd_inp = xp_take_first_n(node_ebd_ext, 1, nloc)
         # repflows
         node_ebd, edge_ebd, h2, rot_mat, sw = self.repflows(
             nlist,

deepmd/dpmodel/descriptor/repflows.py

@@ -13,6 +13,7 @@
 from deepmd.dpmodel.array_api import (
     Array,
     xp_take_along_axis,
+    xp_take_first_n,
 )
 from deepmd.dpmodel.common import (
     to_numpy_array,
@@ -562,7 +563,7 @@ def call(
 
         # get node embedding
         # nb x nloc x tebd_dim
-        atype_embd = atype_embd_ext[:, :nloc, :]
+        atype_embd = xp_take_first_n(atype_embd_ext, 1, nloc)
         assert list(atype_embd.shape) == [nframes, nloc, self.n_dim]
 
         node_ebd = self.act(atype_embd)
@@ -641,7 +642,7 @@ def call(
         angle_ebd = self.angle_embd(angle_input)
 
         # nb x nall x n_dim
-        mapping = xp.tile(xp.reshape(mapping, (nframes, -1, 1)), (1, 1, self.n_dim))
+        mapping = xp.tile(xp.expand_dims(mapping, axis=-1), (1, 1, self.n_dim))
         for idx, ll in enumerate(self.layers):
             # node_ebd:     nb x nloc x n_dim
             # node_ebd_ext: nb x nall x n_dim
@@ -1421,7 +1422,7 @@ def call(
         n_edge = (
             int(xp.sum(xp.astype(nlist_mask, xp.int32))) if self.use_dynamic_sel else 0
         )
-        node_ebd = node_ebd_ext[:, :nloc, :]
+        node_ebd = xp_take_first_n(node_ebd_ext, 1, nloc)
         assert (nb, nloc) == node_ebd.shape[:2]
         if not self.use_dynamic_sel:
             assert (nb, nloc, nnei) == h2.shape[:3]

source/api_cc/tests/test_deeppot_dpa3_pt.cc

@@ -0,0 +1,308 @@
+// SPDX-License-Identifier: LGPL-3.0-or-later
+// Test C++ inference for PT (.pth) backend with DPA3 (mixed-type) model.
+// Reference values generated by source/tests/infer/gen_dpa3.py.
+#include <gtest/gtest.h>
+
+#include <algorithm>
+#include <cmath>
+#include <fstream>
+#include <vector>
+
+#include "DeepPot.h"
+#include "neighbor_list.h"
+#include "test_utils.h"
+
+// DPA3 models need relaxed epsilon
+#undef EPSILON
+#define EPSILON (std::is_same<VALUETYPE, double>::value ? 1e-7 : 1e-1)
+
+template <class VALUETYPE>
+class TestInferDeepPotDpa3Pt : public ::testing::Test {
+ protected:
+  std::vector<VALUETYPE> coord = {12.83, 2.56, 2.18, 12.09, 2.87, 2.74,
+                                  00.25, 3.32, 1.68, 3.36,  3.00, 1.81,
+                                  3.51,  2.51, 2.60, 4.27,  3.22, 1.56};
+  std::vector<int> atype = {0, 1, 1, 0, 1, 1};
+  std::vector<VALUETYPE> box = {13., 0., 0., 0., 13., 0., 0., 0., 13.};
+  // Generated by source/tests/infer/gen_dpa3.py (PBC)
+  // Same weights as deeppot_dpa3.pt2
+  std::vector<VALUETYPE> expected_e = {
+      2.733142942358297023e-01, 2.768815473296480922e-01,
+      2.781664369968356865e-01, 2.697839344989072519e-01,
+      2.741210600049306945e-01, 2.752870928812235496e-01};
+  std::vector<VALUETYPE> expected_f = {
+      -1.962618723134541832e-02, 4.287158582278347702e-02,
+      7.640666386947853050e-03,  5.554130248696588501e-02,
+      -6.501206231527984977e-03, -4.524468847893595158e-02,
+      -3.851051736663693714e-02, -3.620789238677154381e-02,
+      3.756162244251591564e-02,  6.729090678104879264e-02,
+      -2.430710555108604037e-02, 4.496058666120762021e-02,
+      9.285825331084011924e-03,  5.623126339971108029e-02,
+      -8.776072674283137698e-02, -7.398133000111631330e-02,
+      -3.208664505310900028e-02, 4.284253973109593966e-02};
+  std::vector<VALUETYPE> expected_v = {
+      -2.519191242984861884e-02, -7.976296517418629550e-04,
+      2.293255716383547221e-02,  -1.129879902880513709e-04,
+      -2.480533869648754441e-02, 5.147545203263749480e-03,
+      2.250634701911344987e-02,  5.288887046140826331e-03,
+      -2.010244267109611085e-02, -1.779331319768159489e-02,
+      3.093850189397499839e-03,  1.469388965841003300e-02,
+      -3.857294749719837688e-03, 1.122172669801067097e-03,
+      3.015485878866499582e-03,  1.588838841470147090e-02,
+      -2.814760933954751562e-03, -1.277216714527013713e-02,
+      -8.763367643346370306e-03, -1.305889135368112908e-02,
+      1.181350951828694096e-02,  -6.506014073233991855e-03,
+      -6.021216432246893902e-03, 6.406967309407277100e-03,
+      1.054423249710041179e-02,  1.210616766999832172e-02,
+      -1.127472660426425549e-02, -3.873334330831591787e-02,
+      -3.620067664760272686e-03, 1.173198873109224322e-03,
+      -3.979800321914496279e-03, -1.483777776121806245e-02,
+      2.311848485249741111e-02,  1.659292900032220339e-03,
+      2.315104663227764842e-02,  -3.645194750481960122e-02,
+      -1.668107738824501848e-04, -7.331929353596922626e-03,
+      1.141573012886789966e-02,  -1.498650485705460686e-03,
+      -1.339178008942835431e-02, 2.104129816063767672e-02,
+      2.247013447171188061e-03,  2.035538814221872148e-02,
+      -3.195007182084359104e-02, -2.339460083073257798e-02,
+      -1.001949167693141039e-02, 1.320033846426920537e-02,
+      -1.577941189045228843e-02, -6.283307183655661120e-03,
+      8.237968913765561507e-03,  2.238394952866012630e-02,
+      8.881021761757389166e-03,  -1.162377795308391741e-02};
+  int natoms;
+  double expected_tot_e;
+  std::vector<VALUETYPE> expected_tot_v;
+
+  deepmd::DeepPot dp;
+
+  void SetUp() override {
+#ifndef BUILD_PYTORCH
+    GTEST_SKIP() << "Skip because PyTorch support is not enabled.";
+#endif
+    dp.init("../../tests/infer/deeppot_dpa3.pth");
+
+    natoms = expected_e.size();
+    EXPECT_EQ(natoms * 3, expected_f.size());
+    EXPECT_EQ(natoms * 9, expected_v.size());
+    expected_tot_e = 0.;
+    expected_tot_v.resize(9);
+    std::fill(expected_tot_v.begin(), expected_tot_v.end(), 0.);
+    for (int ii = 0; ii < natoms; ++ii) {
+      expected_tot_e += expected_e[ii];
+    }
+    for (int ii = 0; ii < natoms; ++ii) {
+      for (int dd = 0; dd < 9; ++dd) {
+        expected_tot_v[dd] += expected_v[ii * 9 + dd];
+      }
+    }
+  };
+
+  void TearDown() override {};
+};
+
+TYPED_TEST_SUITE(TestInferDeepPotDpa3Pt, ValueTypes);
+
+TYPED_TEST(TestInferDeepPotDpa3Pt, cpu_build_nlist) {
+  using VALUETYPE = TypeParam;
+  std::vector<VALUETYPE>& coord = this->coord;
+  std::vector<int>& atype = this->atype;
+  std::vector<VALUETYPE>& box = this->box;
+  std::vector<VALUETYPE>& expected_f = this->expected_f;
+  int& natoms = this->natoms;
+  double& expected_tot_e = this->expected_tot_e;
+  std::vector<VALUETYPE>& expected_tot_v = this->expected_tot_v;
+  deepmd::DeepPot& dp = this->dp;
+  double ener;
+  std::vector<VALUETYPE> force, virial;
+  dp.compute(ener, force, virial, coord, atype, box);
+
+  EXPECT_EQ(force.size(), natoms * 3);
+  EXPECT_EQ(virial.size(), 9);
+
+  EXPECT_LT(fabs(ener - expected_tot_e), EPSILON);
+  for (int ii = 0; ii < natoms * 3; ++ii) {
+    EXPECT_LT(fabs(force[ii] - expected_f[ii]), EPSILON);
+  }
+  for (int ii = 0; ii < 3 * 3; ++ii) {
+    EXPECT_LT(fabs(virial[ii] - expected_tot_v[ii]), EPSILON);
+  }
+}
+
+TYPED_TEST(TestInferDeepPotDpa3Pt, cpu_build_nlist_atomic) {
+  using VALUETYPE = TypeParam;
+  std::vector<VALUETYPE>& coord = this->coord;
+  std::vector<int>& atype = this->atype;
+  std::vector<VALUETYPE>& box = this->box;
+  std::vector<VALUETYPE>& expected_e = this->expected_e;
+  std::vector<VALUETYPE>& expected_f = this->expected_f;
+  std::vector<VALUETYPE>& expected_v = this->expected_v;
+  int& natoms = this->natoms;
+  double& expected_tot_e = this->expected_tot_e;
+  std::vector<VALUETYPE>& expected_tot_v = this->expected_tot_v;
+  deepmd::DeepPot& dp = this->dp;
+  double ener;
+  std::vector<VALUETYPE> force, virial, atom_ener, atom_vir;
+  dp.compute(ener, force, virial, atom_ener, atom_vir, coord, atype, box);
+
+  EXPECT_EQ(force.size(), natoms * 3);
+  EXPECT_EQ(virial.size(), 9);
+  EXPECT_EQ(atom_ener.size(), natoms);
+  EXPECT_EQ(atom_vir.size(), natoms * 9);
+
+  EXPECT_LT(fabs(ener - expected_tot_e), EPSILON);
+  for (int ii = 0; ii < natoms * 3; ++ii) {
+    EXPECT_LT(fabs(force[ii] - expected_f[ii]), EPSILON);
+  }
+  for (int ii = 0; ii < 3 * 3; ++ii) {
+    EXPECT_LT(fabs(virial[ii] - expected_tot_v[ii]), EPSILON);
+  }
+  for (int ii = 0; ii < natoms; ++ii) {
+    EXPECT_LT(fabs(atom_ener[ii] - expected_e[ii]), EPSILON);
+  }
+  for (int ii = 0; ii < natoms * 9; ++ii) {
+    EXPECT_LT(fabs(atom_vir[ii] - expected_v[ii]), EPSILON);
+  }
+}
+
+template <class VALUETYPE>
+class TestInferDeepPotDpa3PtNoPbc : public ::testing::Test {
+ protected:
+  std::vector<VALUETYPE> coord = {12.83, 2.56, 2.18, 12.09, 2.87, 2.74,
+                                  00.25, 3.32, 1.68, 3.36,  3.00, 1.81,
+                                  3.51,  2.51, 2.60, 4.27,  3.22, 1.56};
+  std::vector<int> atype = {0, 1, 1, 0, 1, 1};
+  std::vector<VALUETYPE> box = {};
+  // Generated by source/tests/infer/gen_dpa3.py (NoPbc)
+  std::vector<VALUETYPE> expected_e = {
+      2.748896667984845887e-01, 2.803947322373078754e-01,
+      2.865499847997139971e-01, 2.695555136277474895e-01,
+      2.739584531066059925e-01, 2.752217127378932537e-01};
+  std::vector<VALUETYPE> expected_f = {
+      -4.469562373941994571e-02, 1.872384237732456838e-02,
+      3.382371526226372882e-02,  4.469562373941994571e-02,
+      -1.872384237732456838e-02, -3.382371526226372882e-02,
+      -8.962417443747255821e-04, 6.973117535150641388e-05,
+      3.708588577163370883e-05,  6.643516471939500678e-02,
+      -2.418189932122343649e-02, 4.484243027251725439e-02,
+      9.031619071676464522e-03,  5.637239343551967569e-02,
+      -8.796029317613156262e-02, -7.457054204669674724e-02,
+      -3.226022528964775371e-02, 4.308077701784267244e-02};
+  std::vector<VALUETYPE> expected_v = {
+      -1.634330450074628072e-02, 6.846519453015231793e-03,
+      1.236790610867266604e-02,  6.846519453015259549e-03,
+      -2.868136527614494058e-03, -5.181149856335852399e-03,
+      1.236790610867266604e-02,  -5.181149856335859338e-03,
+      -9.359496514671244993e-03, -1.673145706642453767e-02,
+      7.009123906204950405e-03,  1.266164318540249911e-02,
+      7.009123906204922649e-03,  -2.936254609356120371e-03,
+      -5.304201874965906727e-03, 1.266164318540247136e-02,
+      -5.304201874965899788e-03, -9.581784032196449807e-03,
+      2.483905957089865488e-03,  -1.710616363479115602e-04,
+      -5.347582359011894028e-05, -1.996686279554130779e-04,
+      1.446275632786597548e-05,  2.638112328458543858e-06,
+      -1.197563523836930226e-04, 1.205600575305949503e-05,
+      -4.593499883389132697e-06, -4.089897480719173473e-02,
+      -3.495830205935246404e-03, 1.154978330068986980e-03,
+      -3.627142383941225900e-03, -1.488475129792680290e-02,
+      2.311785022979555293e-02,  1.347848716528848856e-03,
+      2.315545736893441509e-02,  -3.642400982788428221e-02,
+      -1.119743233540158867e-03, -7.327254171127076110e-03,
+      1.144439607350029517e-02,  -1.403015516843159061e-03,
+      -1.349644754565121341e-02, 2.117430870829728473e-02,
+      2.103115217604090148e-03,  2.047643373328661420e-02,
+      -3.212706064943796069e-02, -2.418232649309504101e-02,
+      -1.012366394018440752e-02, 1.334822742508814941e-02,
+      -1.588798342485496506e-02, -6.330672283764562924e-03,
+      8.295385033255518736e-03,  2.256291842331806241e-02,
+      8.946234975702738179e-03,  -1.170798305154926999e-02};
+  int natoms;
+  double expected_tot_e;
+  std::vector<VALUETYPE> expected_tot_v;
+
+  deepmd::DeepPot dp;
+
+  void SetUp() override {
+#ifndef BUILD_PYTORCH
+    GTEST_SKIP() << "Skip because PyTorch support is not enabled.";
+#endif
+    dp.init("../../tests/infer/deeppot_dpa3.pth");
+
+    natoms = expected_e.size();
+    EXPECT_EQ(natoms * 3, expected_f.size());
+    EXPECT_EQ(natoms * 9, expected_v.size());
+    expected_tot_e = 0.;
+    expected_tot_v.resize(9);
+    std::fill(expected_tot_v.begin(), expected_tot_v.end(), 0.);
+    for (int ii = 0; ii < natoms; ++ii) {
+      expected_tot_e += expected_e[ii];
+    }
+    for (int ii = 0; ii < natoms; ++ii) {
+      for (int dd = 0; dd < 9; ++dd) {
+        expected_tot_v[dd] += expected_v[ii * 9 + dd];
+      }
+    }
+  };
+
+  void TearDown() override {};
+};
+
+TYPED_TEST_SUITE(TestInferDeepPotDpa3PtNoPbc, ValueTypes);
+
+TYPED_TEST(TestInferDeepPotDpa3PtNoPbc, cpu_build_nlist) {
+  using VALUETYPE = TypeParam;
+  std::vector<VALUETYPE>& coord = this->coord;
+  std::vector<int>& atype = this->atype;
+  std::vector<VALUETYPE>& box = this->box;
+  std::vector<VALUETYPE>& expected_f = this->expected_f;
+  int& natoms = this->natoms;
+  double& expected_tot_e = this->expected_tot_e;
+  std::vector<VALUETYPE>& expected_tot_v = this->expected_tot_v;
+  deepmd::DeepPot& dp = this->dp;
+  double ener;
+  std::vector<VALUETYPE> force, virial;
+  dp.compute(ener, force, virial, coord, atype, box);
+
+  EXPECT_EQ(force.size(), natoms * 3);
+  EXPECT_EQ(virial.size(), 9);
+
+  EXPECT_LT(fabs(ener - expected_tot_e), EPSILON);
+  for (int ii = 0; ii < natoms * 3; ++ii) {
+    EXPECT_LT(fabs(force[ii] - expected_f[ii]), EPSILON);
+  }
+  for (int ii = 0; ii < 3 * 3; ++ii) {
+    EXPECT_LT(fabs(virial[ii] - expected_tot_v[ii]), EPSILON);
+  }
+}
+
+TYPED_TEST(TestInferDeepPotDpa3PtNoPbc, cpu_lmp_nlist) {
+  using VALUETYPE = TypeParam;
+  std::vector<VALUETYPE>& coord = this->coord;
+  std::vector<int>& atype = this->atype;
+  std::vector<VALUETYPE>& box = this->box;
+  std::vector<VALUETYPE>& expected_f = this->expected_f;
+  int& natoms = this->natoms;
+  double& expected_tot_e = this->expected_tot_e;
+  std::vector<VALUETYPE>& expected_tot_v = this->expected_tot_v;
+  deepmd::DeepPot& dp = this->dp;
+  double ener;
+  std::vector<VALUETYPE> force, virial;
+
+  std::vector<std::vector<int> > nlist_data = {
+      {1, 2, 3, 4, 5}, {0, 2, 3, 4, 5}, {0, 1, 3, 4, 5},
+      {0, 1, 2, 4, 5}, {0, 1, 2, 3, 5}, {0, 1, 2, 3, 4}};
+  std::vector<int> ilist(natoms), numneigh(natoms);
+  std::vector<int*> firstneigh(natoms);
+  deepmd::InputNlist inlist(natoms, &ilist[0], &numneigh[0], &firstneigh[0]);
+  convert_nlist(inlist, nlist_data);
+  dp.compute(ener, force, virial, coord, atype, box, 0, inlist, 0);
+
+  EXPECT_EQ(force.size(), natoms * 3);
+  EXPECT_EQ(virial.size(), 9);
+
+  EXPECT_LT(fabs(ener - expected_tot_e), EPSILON);
+  for (int ii = 0; ii < natoms * 3; ++ii) {
+    EXPECT_LT(fabs(force[ii] - expected_f[ii]), EPSILON);
+  }
+  for (int ii = 0; ii < 3 * 3; ++ii) {
+    EXPECT_LT(fabs(virial[ii] - expected_tot_v[ii]), EPSILON);
+  }
+}