diff --git a/source/api_cc/tests/test_deeppot_dpa_jax.cc b/source/api_cc/tests/test_deeppot_dpa_jax.cc
new file mode 100644
index 0000000000..bd7f1ee7af
--- /dev/null
+++ b/source/api_cc/tests/test_deeppot_dpa_jax.cc
@@ -0,0 +1,393 @@
+// SPDX-License-Identifier: LGPL-3.0-or-later
+#include <fcntl.h>
+#include <gtest/gtest.h>
+#include <sys/stat.h>
+#include <sys/types.h>
+
+#include <algorithm>
+#include <cmath>
+#include <fstream>
+#include <vector>
+
+#include "DeepPot.h"
+#include "neighbor_list.h"
+#include "test_utils.h"
+
+// 1e-10 cannot pass; unclear bug or not
+#undef EPSILON
+#define EPSILON (std::is_same<VALUETYPE, double>::value ? 1e-7 : 1e-1)
+
+template <class VALUETYPE>
+class TestInferDeepPotDpaJAX : 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 the following Python code:
+  // import numpy as np
+  // from deepmd.infer import DeepPot
+  // coord = np.array([
+  //     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
+  // ]).reshape(1, -1)
+  // atype = np.array([0, 1, 1, 0, 1, 1])
+  // box = np.array([13., 0., 0., 0., 13., 0., 0., 0., 13.]).reshape(1, -1)
+  // dp = DeepPot("deeppot_dpa.savedmodel")
+  // e, f, v, ae, av = dp.eval(coord, box, atype, atomic=True)
+  // np.set_printoptions(precision=16)
+  // print(f"{e.ravel()=} {v.ravel()=} {f.ravel()=} {ae.ravel()=}
+  // {av.ravel()=}")
+
+  std::vector<VALUETYPE> expected_e = {-94.24098099691867,  -187.8049502787117,
+                                       -187.80486052083617, -94.24059525229518,
+                                       -187.80366985846246, -187.8042377490619};
+  std::vector<VALUETYPE> expected_f = {
+      -0.0020150115442053, -0.0133389255924977, -0.0014347177433057,
+      -0.0140757358179293, 0.0031373814221557,  0.0098594354314677,
+      0.004755683505073,   0.0099471082374397,  -0.0080868184532793,
+      -0.0086166721574536, 0.0037803939137322,  -0.0075733131286482,
+      0.0037437603038209,  -0.008452527996008,  0.0134837461840424,
+      0.0162079757106944,  0.0049265700151781,  -0.0062483322902769};
+  std::vector<VALUETYPE> expected_v = {
+      0.0133534319524089,  0.0013445914938337,  -0.0029370551651952,
+      0.0002611806151294,  0.004662662211533,   -0.0002717443796319,
+      -0.0027779798869954, -0.0003277976466339, 0.0018284972283065,
+      0.0085710118978246,  0.0003865036653608,  -0.0057964032875089,
+      -0.0014358330222619, 0.0002912625128908,  0.001212630641674,
+      -0.0050582608957046, -0.0001087907763249, 0.0040068757134429,
+      0.0116736349373084,  0.0007055477968445,  -0.0019544933708784,
+      0.0032997459258512,  0.0037887116116712,  -0.0043140890650835,
+      -0.0034418738401156, -0.0029420616852742, 0.0038219676716965,
+      0.0147134944025738,  0.0005214313829998,  -0.0006524136175906,
+      0.0003656980996363,  0.0010046161607714,  -0.0017279359476254,
+      0.000111127036911,   -0.0017063190420654, 0.0030174567965904,
+      0.0104435705455107,  -0.0008704394438241, 0.0012354202650812,
+      0.0009397615830053,  0.0029105236407293,  -0.0044188897903449,
+      -0.0011461513500477, -0.0045759080125852, 0.0070310883421107,
+      0.0089818851995049,  0.0038819466696704,  -0.005443705549253,
+      0.0025390283635246,  0.0012121502955869,  -0.0016998728971157,
+      -0.0032355117893925, -0.0015590242752438, 0.0021980725909838};
+  int natoms;
+  double expected_tot_e;
+  std::vector<VALUETYPE> expected_tot_v;
+
+  deepmd::DeepPot dp;
+
+  void SetUp() override {
+    dp.init("../../tests/infer/deeppot_dpa.savedmodel");
+
+    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(TestInferDeepPotDpaJAX, ValueTypes);
+
+TYPED_TEST(TestInferDeepPotDpaJAX, 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_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;
+  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(TestInferDeepPotDpaJAX, 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 TestInferDeepPotDpaJAXNopbc : 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 the following Python code:
+  // import numpy as np
+  // from deepmd.infer import DeepPot
+  // coord = np.array([
+  //     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
+  // ]).reshape(1, -1)
+  // atype = np.array([0, 1, 1, 0, 1, 1])
+  // box = None
+  // dp = DeepPot("deeppot_dpa.savedmodel")
+  // e, f, v, ae, av = dp.eval(coord, box, atype, atomic=True)
+  // np.set_printoptions(precision=16)
+  // print(f"{e.ravel()=} {v.ravel()=} {f.ravel()=} {ae.ravel()=}
+  // {av.ravel()=}")
+
+  std::vector<VALUETYPE> expected_e = {
+      -94.24457967995595, -187.81100287606412, -187.81417300904738,
+      -94.24110552426328, -187.80431436838532, -187.80457222464983};
+  std::vector<VALUETYPE> expected_f = {
+      0.0051417559595313,  -0.0021539788479118, -0.0038910585639696,
+      -0.0051417559595313, 0.0021539788479118,  0.0038910585639696,
+      -0.0035470615733886, 0.0003602503965239,  -0.0001895679272905,
+      -0.0117361352793328, 0.0034252835112125,  -0.0071824017939095,
+      0.0005398894945495,  -0.0084330745423862, 0.013284532676939,
+      0.0147433073581718,  0.0046475406346498,  -0.005912562955739};
+  std::vector<VALUETYPE> expected_v = {
+      1.8756488030620411e-03,  -7.8574476885035112e-04, -1.4194099050199721e-03,
+      -7.8574476885031816e-04, 3.2916334911298369e-04,  5.9461766291377029e-04,
+      -1.4194099050199305e-03, 5.9461766291377637e-04,  1.0741480362313257e-03,
+      1.9292506069911106e-03,  -8.0819957860435269e-04, -1.4599734323175548e-03,
+      -8.0819957860438912e-04, 3.3857009373966171e-04,  6.1161049191681572e-04,
+      -1.4599734323175878e-03, 6.1161049191681138e-04,  1.1048447595916697e-03,
+      7.6085550162737362e-03,  -9.6772620145935649e-04, 8.3466357433496438e-04,
+      -6.0870001592646837e-04, 7.8216372819095133e-05,  -6.6454748650881927e-05,
+      1.6484218595781155e-04,  -1.4860654308892986e-05, 9.1346879754941710e-06,
+      7.3496442591020616e-03,  -4.4351638231197171e-04, 1.0048648094020743e-03,
+      -6.7678341212752506e-04, 1.0347359647873023e-03,  -1.6645682233463639e-03,
+      1.4376416549492857e-03,  -1.7116654329527997e-03, 2.8516307661164836e-03,
+      1.7683913600324828e-03,  -2.6945055858769765e-04, 3.0381043714224000e-04,
+      9.9317892494769217e-04,  2.9343081937224687e-03,  -4.4381349285303948e-03,
+      -1.6684804850477311e-03, -4.5760063242183471e-03, 6.9407432349694147e-03,
+      7.8021639779489171e-03,  3.5246152053432442e-03,  -4.9415349551876243e-03,
+      2.1362265660905197e-03,  1.2226850609509659e-03,  -1.7016886981679963e-03,
+      -2.7321994901677166e-03, -1.5683141872155984e-03, 2.1960566952029213e-03};
+  int natoms;
+  double expected_tot_e;
+  std::vector<VALUETYPE> expected_tot_v;
+
+  deepmd::DeepPot dp;
+
+  void SetUp() override {
+    dp.init("../../tests/infer/deeppot_dpa.savedmodel");
+
+    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(TestInferDeepPotDpaJAXNopbc, ValueTypes);
+
+TYPED_TEST(TestInferDeepPotDpaJAXNopbc, 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_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;
+  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(TestInferDeepPotDpaJAXNopbc, 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);
+  }
+}
+
+TYPED_TEST(TestInferDeepPotDpaJAXNopbc, 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_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;
+
+  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);
+  }
+}
+
+TYPED_TEST(TestInferDeepPotDpaJAXNopbc, cpu_lmp_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;
+
+  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, atom_ener, atom_vir, coord, atype, box, 0,
+             inlist, 0);
+
+  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);
+  }
+}