From 387708026d843ce568d01885f32f61d13d07bc81 Mon Sep 17 00:00:00 2001
From: "njzjz-bot (driven by OpenClaw (model:
 custom-chat-jinzhezeng-group/gpt-5.5))[bot]"
 <48687836+njzjz-bot@users.noreply.github.com>
Date: Fri, 29 May 2026 18:49:01 +0000
Subject: [PATCH 1/8] fix(pt): clone inference coords before enabling grad

Avoid enabling gradients in-place on LAMMPS-provided coordinate views.
Clone the extended coordinates before building force graphs and pass that
cloned tensor to derivative generation while keeping auxiliary metadata
private to model output conversion.

Closes #5165

Authored by OpenClaw (model: custom-chat-jinzhezeng-group/gpt-5.5)
---
 .../model/atomic_model/base_atomic_model.py   |  4 ++-
 .../pt/model/atomic_model/dp_atomic_model.py  |  4 ++-
 .../model/atomic_model/linear_atomic_model.py |  4 ++-
 .../atomic_model/pairtab_atomic_model.py      |  5 ++-
 deepmd/pt/model/model/make_model.py           |  5 +--
 source/tests/pt/model/test_dp_atomic_model.py | 28 +++++++++++++++
 source/tests/pt/model/test_dp_model.py        | 35 +++++++++++++++++++
 7 files changed, 79 insertions(+), 6 deletions(-)

diff --git a/deepmd/pt/model/atomic_model/base_atomic_model.py b/deepmd/pt/model/atomic_model/base_atomic_model.py
index 8605db9359..4ab995004f 100644
--- a/deepmd/pt/model/atomic_model/base_atomic_model.py
+++ b/deepmd/pt/model/atomic_model/base_atomic_model.py
@@ -383,7 +383,9 @@ def forward_common_atomic(
         if self.atom_excl is not None:
             atom_mask *= self.atom_excl(atype)
 
-        for kk in ret_dict.keys():
+        for kk in list(ret_dict.keys()):
+            if kk.startswith("_"):
+                continue
             out_shape = ret_dict[kk].shape
             out_shape2 = 1
             for ss in out_shape[2:]:
diff --git a/deepmd/pt/model/atomic_model/dp_atomic_model.py b/deepmd/pt/model/atomic_model/dp_atomic_model.py
index 783ee9e766..f7418948a8 100644
--- a/deepmd/pt/model/atomic_model/dp_atomic_model.py
+++ b/deepmd/pt/model/atomic_model/dp_atomic_model.py
@@ -272,7 +272,7 @@ def forward_atomic(
         nframes, nloc, nnei = nlist.shape
         atype = extended_atype[:, :nloc]
         if self.do_grad_r() or self.do_grad_c():
-            extended_coord.requires_grad_(True)
+            extended_coord = extended_coord.detach().clone().requires_grad_(True)
 
         # Handle default chg_spin if descriptor supports it
         if self.add_chg_spin_ebd and charge_spin is None:
@@ -302,6 +302,8 @@ def forward_atomic(
             fparam=fparam,
             aparam=aparam,
         )
+        if self.do_grad_r() or self.do_grad_c():
+            fit_ret["_force_coord"] = extended_coord
         if self.enable_eval_fitting_last_layer_hook:
             assert "middle_output" in fit_ret, (
                 "eval_fitting_last_layer not supported for this fitting net!"
diff --git a/deepmd/pt/model/atomic_model/linear_atomic_model.py b/deepmd/pt/model/atomic_model/linear_atomic_model.py
index 5c0f616634..8d57a802d5 100644
--- a/deepmd/pt/model/atomic_model/linear_atomic_model.py
+++ b/deepmd/pt/model/atomic_model/linear_atomic_model.py
@@ -259,7 +259,7 @@ def forward_atomic(
         """
         nframes, nloc, nnei = nlist.shape
         if self.do_grad_r() or self.do_grad_c():
-            extended_coord.requires_grad_(True)
+            extended_coord = extended_coord.detach().clone().requires_grad_(True)
         extended_coord = extended_coord.view(nframes, -1, 3)
         sorted_rcuts, sorted_sels = self._sort_rcuts_sels()
         nlists = build_multiple_neighbor_list(
@@ -304,6 +304,8 @@ def forward_atomic(
                 dim=0,
             ),
         }  # (nframes, nloc, 1)
+        if self.do_grad_r() or self.do_grad_c():
+            fit_ret["_force_coord"] = extended_coord
         return fit_ret
 
     def apply_out_stat(
diff --git a/deepmd/pt/model/atomic_model/pairtab_atomic_model.py b/deepmd/pt/model/atomic_model/pairtab_atomic_model.py
index 5750f7cfd1..6b7249d9a3 100644
--- a/deepmd/pt/model/atomic_model/pairtab_atomic_model.py
+++ b/deepmd/pt/model/atomic_model/pairtab_atomic_model.py
@@ -276,7 +276,7 @@ def forward_atomic(
         nframes, nloc, nnei = nlist.shape
         extended_coord = extended_coord.view(nframes, -1, 3)
         if self.do_grad_r() or self.do_grad_c():
-            extended_coord.requires_grad_(True)
+            extended_coord = extended_coord.detach().clone().requires_grad_(True)
 
         # this will mask all -1 in the nlist
         mask = nlist >= 0
@@ -313,6 +313,9 @@ def forward_atomic(
             dim=-1,
         ).unsqueeze(-1)
 
+        if self.do_grad_r() or self.do_grad_c():
+            atomic_energy = atomic_energy + 0.0 * extended_coord.sum()[..., None, None]
+            return {"energy": atomic_energy, "_force_coord": extended_coord}
         return {"energy": atomic_energy}
 
     def _pair_tabulated_inter(
diff --git a/deepmd/pt/model/model/make_model.py b/deepmd/pt/model/model/make_model.py
index 713eab3d8c..ea1df9bb02 100644
--- a/deepmd/pt/model/model/make_model.py
+++ b/deepmd/pt/model/model/make_model.py
@@ -305,7 +305,7 @@ def forward_common_lower(
             cc_ext, _, fp, ap, input_prec = self._input_type_cast(
                 extended_coord, fparam=fparam, aparam=aparam
             )
-            del extended_coord, fparam, aparam
+            del fparam, aparam
             atomic_ret = self.atomic_model.forward_common_atomic(
                 cc_ext,
                 extended_atype,
@@ -316,10 +316,11 @@ def forward_common_lower(
                 comm_dict=comm_dict,
                 charge_spin=charge_spin,
             )
+            force_coord = atomic_ret.pop("_force_coord", cc_ext)
             model_predict = fit_output_to_model_output(
                 atomic_ret,
                 self.atomic_output_def(),
-                cc_ext,
+                force_coord,
                 do_atomic_virial=do_atomic_virial,
                 create_graph=self.training,
                 mask=atomic_ret["mask"] if "mask" in atomic_ret else None,
diff --git a/source/tests/pt/model/test_dp_atomic_model.py b/source/tests/pt/model/test_dp_atomic_model.py
index 6d6a22f357..d77c986bb4 100644
--- a/source/tests/pt/model/test_dp_atomic_model.py
+++ b/source/tests/pt/model/test_dp_atomic_model.py
@@ -74,6 +74,34 @@ def test_self_consistency(self) -> None:
                 to_numpy_array(ret1["energy"]),
             )
 
+    def test_forward_common_atomic_accepts_leaf_view_input(self) -> None:
+        ds = DescrptSeA(
+            self.rcut,
+            self.rcut_smth,
+            self.sel,
+        ).to(env.DEVICE)
+        ft = InvarFitting(
+            "energy",
+            self.nt,
+            ds.get_dim_out(),
+            1,
+            mixed_types=ds.mixed_types(),
+        ).to(env.DEVICE)
+        md0 = DPAtomicModel(ds, ft, type_map=["foo", "bar"]).to(env.DEVICE)
+
+        coord = to_torch_tensor(self.coord_ext).requires_grad_(True)
+        coord_view = coord.view(self.nf, self.nall, 3)
+        args = [
+            coord_view,
+            to_torch_tensor(self.atype_ext),
+            to_torch_tensor(self.nlist),
+        ]
+        atomic_ret = md0.forward_atomic(*args)
+        ret = md0.forward_common_atomic(*args)
+
+        self.assertIn("_force_coord", atomic_ret)
+        self.assertIn("energy", ret)
+
     def test_dp_consistency(self) -> None:
         nf, nloc, nnei = self.nlist.shape
         ds = DPDescrptSeA(
diff --git a/source/tests/pt/model/test_dp_model.py b/source/tests/pt/model/test_dp_model.py
index f4e350869a..a5adeada13 100644
--- a/source/tests/pt/model/test_dp_model.py
+++ b/source/tests/pt/model/test_dp_model.py
@@ -114,6 +114,41 @@ def test_self_consistency(self) -> None:
             atol=self.atol,
         )
 
+    def test_forward_lower_accepts_leaf_view_input(self) -> None:
+        ds = DescrptSeA(
+            self.rcut,
+            self.rcut_smth,
+            self.sel,
+        ).to(env.DEVICE)
+        ft = EnergyFittingNet(
+            self.nt,
+            ds.get_dim_out(),
+            mixed_types=ds.mixed_types(),
+        ).to(env.DEVICE)
+        type_map = ["foo", "bar"]
+        md0 = EnergyModel(ds, ft, type_map=type_map).to(env.DEVICE)
+
+        coord_ext, atype_ext, _ = extend_coord_with_ghosts(
+            to_torch_tensor(self.coord),
+            to_torch_tensor(self.atype),
+            to_torch_tensor(self.cell),
+            self.rcut,
+        )
+        nlist = build_neighbor_list(
+            coord_ext,
+            atype_ext,
+            self.nloc,
+            self.rcut,
+            self.sel,
+            distinguish_types=(not md0.mixed_types()),
+        )
+        coord_view = coord_ext.requires_grad_(True).view(self.nf, -1, 3)
+
+        ret = md0.forward_lower(coord_view, atype_ext, nlist, do_atomic_virial=True)
+
+        self.assertIn("extended_force", ret)
+        self.assertIn("virial", ret)
+
     def test_dp_consistency(self) -> None:
         nf, nloc = self.atype.shape
         nfp, nap = 2, 3

From a101808efb159b4d10d4e0bdb1c87a786686fcd2 Mon Sep 17 00:00:00 2001
From: "njzjz-bot (driven by OpenClaw (model:
 custom-chat-jinzhezeng-group/gpt-5.5))[bot]"
 <48687836+njzjz-bot@users.noreply.github.com>
Date: Sat, 30 May 2026 02:50:43 +0000
Subject: [PATCH 2/8] fix(pt): preserve existing coordinate grad graph

Only clone extended coordinates when gradients are not already enabled.
This keeps existing autodiff and Hessian graphs intact while still avoiding
in-place requires_grad_ on non-grad coordinate views.

Authored by OpenClaw (model: custom-chat-jinzhezeng-group/gpt-5.5)
---
 deepmd/pt/model/atomic_model/dp_atomic_model.py      | 2 +-
 deepmd/pt/model/atomic_model/linear_atomic_model.py  | 2 +-
 deepmd/pt/model/atomic_model/pairtab_atomic_model.py | 2 +-
 3 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/deepmd/pt/model/atomic_model/dp_atomic_model.py b/deepmd/pt/model/atomic_model/dp_atomic_model.py
index f7418948a8..9955a9e09b 100644
--- a/deepmd/pt/model/atomic_model/dp_atomic_model.py
+++ b/deepmd/pt/model/atomic_model/dp_atomic_model.py
@@ -271,7 +271,7 @@ def forward_atomic(
         """
         nframes, nloc, nnei = nlist.shape
         atype = extended_atype[:, :nloc]
-        if self.do_grad_r() or self.do_grad_c():
+        if (self.do_grad_r() or self.do_grad_c()) and not extended_coord.requires_grad:
             extended_coord = extended_coord.detach().clone().requires_grad_(True)
 
         # Handle default chg_spin if descriptor supports it
diff --git a/deepmd/pt/model/atomic_model/linear_atomic_model.py b/deepmd/pt/model/atomic_model/linear_atomic_model.py
index 8d57a802d5..655dad6e9a 100644
--- a/deepmd/pt/model/atomic_model/linear_atomic_model.py
+++ b/deepmd/pt/model/atomic_model/linear_atomic_model.py
@@ -258,7 +258,7 @@ def forward_atomic(
             the result dict, defined by the fitting net output def.
         """
         nframes, nloc, nnei = nlist.shape
-        if self.do_grad_r() or self.do_grad_c():
+        if (self.do_grad_r() or self.do_grad_c()) and not extended_coord.requires_grad:
             extended_coord = extended_coord.detach().clone().requires_grad_(True)
         extended_coord = extended_coord.view(nframes, -1, 3)
         sorted_rcuts, sorted_sels = self._sort_rcuts_sels()
diff --git a/deepmd/pt/model/atomic_model/pairtab_atomic_model.py b/deepmd/pt/model/atomic_model/pairtab_atomic_model.py
index 6b7249d9a3..e1ffe59ea3 100644
--- a/deepmd/pt/model/atomic_model/pairtab_atomic_model.py
+++ b/deepmd/pt/model/atomic_model/pairtab_atomic_model.py
@@ -275,7 +275,7 @@ def forward_atomic(
     ) -> dict[str, torch.Tensor]:
         nframes, nloc, nnei = nlist.shape
         extended_coord = extended_coord.view(nframes, -1, 3)
-        if self.do_grad_r() or self.do_grad_c():
+        if (self.do_grad_r() or self.do_grad_c()) and not extended_coord.requires_grad:
             extended_coord = extended_coord.detach().clone().requires_grad_(True)
 
         # this will mask all -1 in the nlist

From d5270996b5e3267d975e212da259f81af8da0103 Mon Sep 17 00:00:00 2001
From: "njzjz-bot (driven by OpenClaw (model:
 custom-chat-jinzhezeng-group/gpt-5.5))[bot]"
 <48687836+njzjz-bot@users.noreply.github.com>
Date: Sat, 30 May 2026 06:15:24 +0000
Subject: [PATCH 3/8] test(pt): cover grad enabling on coordinate views

---
 source/tests/pt/model/test_dp_atomic_model.py | 2 +-
 source/tests/pt/model/test_dp_model.py        | 3 ++-
 2 files changed, 3 insertions(+), 2 deletions(-)

diff --git a/source/tests/pt/model/test_dp_atomic_model.py b/source/tests/pt/model/test_dp_atomic_model.py
index d77c986bb4..c30b937487 100644
--- a/source/tests/pt/model/test_dp_atomic_model.py
+++ b/source/tests/pt/model/test_dp_atomic_model.py
@@ -89,7 +89,7 @@ def test_forward_common_atomic_accepts_leaf_view_input(self) -> None:
         ).to(env.DEVICE)
         md0 = DPAtomicModel(ds, ft, type_map=["foo", "bar"]).to(env.DEVICE)
 
-        coord = to_torch_tensor(self.coord_ext).requires_grad_(True)
+        coord = to_torch_tensor(self.coord_ext)
         coord_view = coord.view(self.nf, self.nall, 3)
         args = [
             coord_view,
diff --git a/source/tests/pt/model/test_dp_model.py b/source/tests/pt/model/test_dp_model.py
index a5adeada13..37f0fb9caa 100644
--- a/source/tests/pt/model/test_dp_model.py
+++ b/source/tests/pt/model/test_dp_model.py
@@ -142,10 +142,11 @@ def test_forward_lower_accepts_leaf_view_input(self) -> None:
             self.sel,
             distinguish_types=(not md0.mixed_types()),
         )
-        coord_view = coord_ext.requires_grad_(True).view(self.nf, -1, 3)
+        coord_view = coord_ext.view(self.nf, -1, 3)
 
         ret = md0.forward_lower(coord_view, atype_ext, nlist, do_atomic_virial=True)
 
+        self.assertFalse(coord_view.requires_grad)
         self.assertIn("extended_force", ret)
         self.assertIn("virial", ret)
 

From 889eff4b86694dcc1249a09a3b9f8ca933fe280b Mon Sep 17 00:00:00 2001
From: "njzjz-bot (driven by OpenClaw (model:
 custom-chat-jinzhezeng-group/gpt-5.5))[bot]"
 <48687836+njzjz-bot@users.noreply.github.com>
Date: Sat, 30 May 2026 10:03:53 +0000
Subject: [PATCH 4/8] fix(pt): prepare force coordinate in model layer

---
 deepmd/pt/model/atomic_model/base_atomic_model.py    | 4 +---
 deepmd/pt/model/atomic_model/dp_atomic_model.py      | 4 +---
 deepmd/pt/model/atomic_model/linear_atomic_model.py  | 4 +---
 deepmd/pt/model/atomic_model/pairtab_atomic_model.py | 5 +----
 deepmd/pt/model/model/make_model.py                  | 9 ++++++---
 source/tests/pt/model/test_dp_atomic_model.py        | 2 --
 6 files changed, 10 insertions(+), 18 deletions(-)

diff --git a/deepmd/pt/model/atomic_model/base_atomic_model.py b/deepmd/pt/model/atomic_model/base_atomic_model.py
index 4ab995004f..8605db9359 100644
--- a/deepmd/pt/model/atomic_model/base_atomic_model.py
+++ b/deepmd/pt/model/atomic_model/base_atomic_model.py
@@ -383,9 +383,7 @@ def forward_common_atomic(
         if self.atom_excl is not None:
             atom_mask *= self.atom_excl(atype)
 
-        for kk in list(ret_dict.keys()):
-            if kk.startswith("_"):
-                continue
+        for kk in ret_dict.keys():
             out_shape = ret_dict[kk].shape
             out_shape2 = 1
             for ss in out_shape[2:]:
diff --git a/deepmd/pt/model/atomic_model/dp_atomic_model.py b/deepmd/pt/model/atomic_model/dp_atomic_model.py
index 9955a9e09b..25821fd945 100644
--- a/deepmd/pt/model/atomic_model/dp_atomic_model.py
+++ b/deepmd/pt/model/atomic_model/dp_atomic_model.py
@@ -272,7 +272,7 @@ def forward_atomic(
         nframes, nloc, nnei = nlist.shape
         atype = extended_atype[:, :nloc]
         if (self.do_grad_r() or self.do_grad_c()) and not extended_coord.requires_grad:
-            extended_coord = extended_coord.detach().clone().requires_grad_(True)
+            extended_coord.requires_grad_(True)
 
         # Handle default chg_spin if descriptor supports it
         if self.add_chg_spin_ebd and charge_spin is None:
@@ -302,8 +302,6 @@ def forward_atomic(
             fparam=fparam,
             aparam=aparam,
         )
-        if self.do_grad_r() or self.do_grad_c():
-            fit_ret["_force_coord"] = extended_coord
         if self.enable_eval_fitting_last_layer_hook:
             assert "middle_output" in fit_ret, (
                 "eval_fitting_last_layer not supported for this fitting net!"
diff --git a/deepmd/pt/model/atomic_model/linear_atomic_model.py b/deepmd/pt/model/atomic_model/linear_atomic_model.py
index 655dad6e9a..41fd49f096 100644
--- a/deepmd/pt/model/atomic_model/linear_atomic_model.py
+++ b/deepmd/pt/model/atomic_model/linear_atomic_model.py
@@ -259,7 +259,7 @@ def forward_atomic(
         """
         nframes, nloc, nnei = nlist.shape
         if (self.do_grad_r() or self.do_grad_c()) and not extended_coord.requires_grad:
-            extended_coord = extended_coord.detach().clone().requires_grad_(True)
+            extended_coord.requires_grad_(True)
         extended_coord = extended_coord.view(nframes, -1, 3)
         sorted_rcuts, sorted_sels = self._sort_rcuts_sels()
         nlists = build_multiple_neighbor_list(
@@ -304,8 +304,6 @@ def forward_atomic(
                 dim=0,
             ),
         }  # (nframes, nloc, 1)
-        if self.do_grad_r() or self.do_grad_c():
-            fit_ret["_force_coord"] = extended_coord
         return fit_ret
 
     def apply_out_stat(
diff --git a/deepmd/pt/model/atomic_model/pairtab_atomic_model.py b/deepmd/pt/model/atomic_model/pairtab_atomic_model.py
index e1ffe59ea3..9277725c23 100644
--- a/deepmd/pt/model/atomic_model/pairtab_atomic_model.py
+++ b/deepmd/pt/model/atomic_model/pairtab_atomic_model.py
@@ -276,7 +276,7 @@ def forward_atomic(
         nframes, nloc, nnei = nlist.shape
         extended_coord = extended_coord.view(nframes, -1, 3)
         if (self.do_grad_r() or self.do_grad_c()) and not extended_coord.requires_grad:
-            extended_coord = extended_coord.detach().clone().requires_grad_(True)
+            extended_coord.requires_grad_(True)
 
         # this will mask all -1 in the nlist
         mask = nlist >= 0
@@ -313,9 +313,6 @@ def forward_atomic(
             dim=-1,
         ).unsqueeze(-1)
 
-        if self.do_grad_r() or self.do_grad_c():
-            atomic_energy = atomic_energy + 0.0 * extended_coord.sum()[..., None, None]
-            return {"energy": atomic_energy, "_force_coord": extended_coord}
         return {"energy": atomic_energy}
 
     def _pair_tabulated_inter(
diff --git a/deepmd/pt/model/model/make_model.py b/deepmd/pt/model/model/make_model.py
index ea1df9bb02..78705b153c 100644
--- a/deepmd/pt/model/model/make_model.py
+++ b/deepmd/pt/model/model/make_model.py
@@ -305,9 +305,13 @@ def forward_common_lower(
             cc_ext, _, fp, ap, input_prec = self._input_type_cast(
                 extended_coord, fparam=fparam, aparam=aparam
             )
-            del fparam, aparam
+            del extended_coord, fparam, aparam
+            force_coord = cc_ext
+            if self.atomic_model.do_grad_r() or self.atomic_model.do_grad_c():
+                if not force_coord.requires_grad:
+                    force_coord = force_coord.clone().requires_grad_(True)
             atomic_ret = self.atomic_model.forward_common_atomic(
-                cc_ext,
+                force_coord,
                 extended_atype,
                 nlist,
                 mapping=mapping,
@@ -316,7 +320,6 @@ def forward_common_lower(
                 comm_dict=comm_dict,
                 charge_spin=charge_spin,
             )
-            force_coord = atomic_ret.pop("_force_coord", cc_ext)
             model_predict = fit_output_to_model_output(
                 atomic_ret,
                 self.atomic_output_def(),
diff --git a/source/tests/pt/model/test_dp_atomic_model.py b/source/tests/pt/model/test_dp_atomic_model.py
index c30b937487..ddc8f6d5ed 100644
--- a/source/tests/pt/model/test_dp_atomic_model.py
+++ b/source/tests/pt/model/test_dp_atomic_model.py
@@ -96,10 +96,8 @@ def test_forward_common_atomic_accepts_leaf_view_input(self) -> None:
             to_torch_tensor(self.atype_ext),
             to_torch_tensor(self.nlist),
         ]
-        atomic_ret = md0.forward_atomic(*args)
         ret = md0.forward_common_atomic(*args)
 
-        self.assertIn("_force_coord", atomic_ret)
         self.assertIn("energy", ret)
 
     def test_dp_consistency(self) -> None:

From 8e1670d9988a8ecaf1cd9aabc1b4322b800b6401 Mon Sep 17 00:00:00 2001
From: "njzjz-bot (driven by OpenClaw (model:
 custom-chat-jinzhezeng-group/gpt-5.5))[bot]"
 <48687836+njzjz-bot@users.noreply.github.com>
Date: Sat, 30 May 2026 10:48:26 +0000
Subject: [PATCH 5/8] fix(pt): clone DP atomic coords before enabling grad

---
 deepmd/pt/model/atomic_model/dp_atomic_model.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/deepmd/pt/model/atomic_model/dp_atomic_model.py b/deepmd/pt/model/atomic_model/dp_atomic_model.py
index 25821fd945..d59d518cab 100644
--- a/deepmd/pt/model/atomic_model/dp_atomic_model.py
+++ b/deepmd/pt/model/atomic_model/dp_atomic_model.py
@@ -272,7 +272,7 @@ def forward_atomic(
         nframes, nloc, nnei = nlist.shape
         atype = extended_atype[:, :nloc]
         if (self.do_grad_r() or self.do_grad_c()) and not extended_coord.requires_grad:
-            extended_coord.requires_grad_(True)
+            extended_coord = extended_coord.clone().requires_grad_(True)
 
         # Handle default chg_spin if descriptor supports it
         if self.add_chg_spin_ebd and charge_spin is None:

From d8292377d855acc708492badda6a0dd815dd74a6 Mon Sep 17 00:00:00 2001
From: "njzjz-bot (driven by OpenClaw (model:
 custom-chat-jinzhezeng-group/gpt-5.5))[bot]"
 <48687836+njzjz-bot@users.noreply.github.com>
Date: Mon, 1 Jun 2026 12:05:22 +0000
Subject: [PATCH 6/8] test(pt): cover grad-enabled atomic coordinate inputs

Address review feedback for the coordinate gradient regression by cloning PairTab coordinates before reshaping, asserting direct atomic callers do not mutate view inputs, and covering already grad-enabled inputs for DP, model lower, linear, and pairtab paths.

Authored by OpenClaw (model: custom-chat-jinzhezeng-group/gpt-5.5)
---
 .../atomic_model/pairtab_atomic_model.py      |  4 +--
 source/tests/pt/model/test_dp_atomic_model.py | 35 +++++++++++++++++++
 source/tests/pt/model/test_dp_model.py        |  7 ++++
 .../pt/model/test_linear_atomic_model.py      | 11 ++++++
 .../pt/model/test_pairtab_atomic_model.py     |  8 +++++
 5 files changed, 63 insertions(+), 2 deletions(-)

diff --git a/deepmd/pt/model/atomic_model/pairtab_atomic_model.py b/deepmd/pt/model/atomic_model/pairtab_atomic_model.py
index 9277725c23..eb70cc0e78 100644
--- a/deepmd/pt/model/atomic_model/pairtab_atomic_model.py
+++ b/deepmd/pt/model/atomic_model/pairtab_atomic_model.py
@@ -274,9 +274,9 @@ def forward_atomic(
         charge_spin: torch.Tensor | None = None,
     ) -> dict[str, torch.Tensor]:
         nframes, nloc, nnei = nlist.shape
-        extended_coord = extended_coord.view(nframes, -1, 3)
         if (self.do_grad_r() or self.do_grad_c()) and not extended_coord.requires_grad:
-            extended_coord.requires_grad_(True)
+            extended_coord = extended_coord.clone().requires_grad_(True)
+        extended_coord = extended_coord.view(nframes, -1, 3)
 
         # this will mask all -1 in the nlist
         mask = nlist >= 0
diff --git a/source/tests/pt/model/test_dp_atomic_model.py b/source/tests/pt/model/test_dp_atomic_model.py
index ddc8f6d5ed..aa6640992e 100644
--- a/source/tests/pt/model/test_dp_atomic_model.py
+++ b/source/tests/pt/model/test_dp_atomic_model.py
@@ -91,6 +91,9 @@ def test_forward_common_atomic_accepts_leaf_view_input(self) -> None:
 
         coord = to_torch_tensor(self.coord_ext)
         coord_view = coord.view(self.nf, self.nall, 3)
+        coord_view_before = coord_view.detach().clone()
+        self.assertTrue(coord.is_leaf)
+        self.assertTrue(coord_view._is_view())
         args = [
             coord_view,
             to_torch_tensor(self.atype_ext),
@@ -98,7 +101,39 @@ def test_forward_common_atomic_accepts_leaf_view_input(self) -> None:
         ]
         ret = md0.forward_common_atomic(*args)
 
+        self.assertFalse(coord_view.requires_grad)
+        torch.testing.assert_close(coord_view, coord_view_before)
         self.assertIn("energy", ret)
+        self.assertEqual(ret["energy"].shape, (self.nf, self.nloc, 1))
+        self.assertTrue(torch.isfinite(ret["energy"]).all())
+
+    def test_forward_common_atomic_preserves_grad_enabled_input(self) -> None:
+        ds = DescrptSeA(
+            self.rcut,
+            self.rcut_smth,
+            self.sel,
+        ).to(env.DEVICE)
+        ft = InvarFitting(
+            "energy",
+            self.nt,
+            ds.get_dim_out(),
+            1,
+            mixed_types=ds.mixed_types(),
+        ).to(env.DEVICE)
+        md0 = DPAtomicModel(ds, ft, type_map=["foo", "bar"]).to(env.DEVICE)
+
+        coord = to_torch_tensor(self.coord_ext)
+        coord_view = coord.view(self.nf, self.nall, 3).clone().requires_grad_(True)
+        args = [
+            coord_view,
+            to_torch_tensor(self.atype_ext),
+            to_torch_tensor(self.nlist),
+        ]
+        ret = md0.forward_common_atomic(*args)
+        ret["energy"].sum().backward()
+
+        self.assertTrue(coord_view.requires_grad)
+        self.assertIsNotNone(coord_view.grad)
 
     def test_dp_consistency(self) -> None:
         nf, nloc, nnei = self.nlist.shape
diff --git a/source/tests/pt/model/test_dp_model.py b/source/tests/pt/model/test_dp_model.py
index 37f0fb9caa..63a5b45f34 100644
--- a/source/tests/pt/model/test_dp_model.py
+++ b/source/tests/pt/model/test_dp_model.py
@@ -150,6 +150,13 @@ def test_forward_lower_accepts_leaf_view_input(self) -> None:
         self.assertIn("extended_force", ret)
         self.assertIn("virial", ret)
 
+        coord_view_grad = coord_ext.view(self.nf, -1, 3).clone().requires_grad_(True)
+        ret = md0.forward_lower(coord_view_grad, atype_ext, nlist, do_atomic_virial=True)
+
+        self.assertTrue(coord_view_grad.requires_grad)
+        self.assertIn("extended_force", ret)
+        self.assertIn("virial", ret)
+
     def test_dp_consistency(self) -> None:
         nf, nloc = self.atype.shape
         nfp, nap = 2, 3
diff --git a/source/tests/pt/model/test_linear_atomic_model.py b/source/tests/pt/model/test_linear_atomic_model.py
index 038bb04a5a..6e58159ad4 100644
--- a/source/tests/pt/model/test_linear_atomic_model.py
+++ b/source/tests/pt/model/test_linear_atomic_model.py
@@ -181,6 +181,17 @@ def test_self_consistency(self) -> None:
             to_numpy_array(ret0["energy"]), ret2["energy"], atol=0.001, rtol=0.001
         )
 
+    def test_forward_atomic_preserves_grad_enabled_input(self) -> None:
+        args = [
+            to_torch_tensor(ii) for ii in [self.coord_ext, self.atype_ext, self.nlist]
+        ]
+        args[0] = args[0].view(self.nf, self.nall, 3).clone().requires_grad_(True)
+        ret = self.md0.forward_atomic(*args)
+        ret["energy"].sum().backward()
+
+        self.assertTrue(args[0].requires_grad)
+        self.assertIsNotNone(args[0].grad)
+
     def test_jit(self) -> None:
         md1 = torch.jit.script(self.md1)
         # atomic model no more export methods
diff --git a/source/tests/pt/model/test_pairtab_atomic_model.py b/source/tests/pt/model/test_pairtab_atomic_model.py
index 0f324cbf51..6c3c325a3e 100644
--- a/source/tests/pt/model/test_pairtab_atomic_model.py
+++ b/source/tests/pt/model/test_pairtab_atomic_model.py
@@ -96,6 +96,14 @@ def test_with_mask(self) -> None:
             result["energy"], expected_result, rtol=0.0001, atol=0.0001
         )
 
+    def test_forward_common_atomic_preserves_grad_enabled_input(self) -> None:
+        coord = self.extended_coord.clone().requires_grad_(True)
+        ret = self.model.forward_atomic(coord, self.extended_atype, self.nlist)
+        ret["energy"].sum().backward()
+
+        self.assertTrue(coord.requires_grad)
+        self.assertIsNotNone(coord.grad)
+
     def test_jit(self) -> None:
         model = torch.jit.script(self.model)
         # atomic model no more export methods

From 912ea4010745ae9346f41117df9c6990368d7adf Mon Sep 17 00:00:00 2001
From: "njzjz-bot (driven by OpenClaw (model:
 custom-chat-jinzhezeng-group/gpt-5.5))[bot]"
 <48687836+njzjz-bot@users.noreply.github.com>
Date: Tue, 2 Jun 2026 04:18:07 +0000
Subject: [PATCH 7/8] fix(pt): clone linear atomic coords before enabling grad

Apply the same clone-before-requires-grad treatment to the linear atomic wrapper so leaf views are not mutated in-place before reshaping.

Add regression coverage for leaf-view inputs in linear and pair-tab atomic models.

Authored by OpenClaw (model: custom-chat-jinzhezeng-group/gpt-5.5)
---
 .../pt/model/atomic_model/linear_atomic_model.py |  2 +-
 .../tests/pt/model/test_linear_atomic_model.py   | 16 ++++++++++++++++
 .../tests/pt/model/test_pairtab_atomic_model.py  | 12 ++++++++++++
 3 files changed, 29 insertions(+), 1 deletion(-)

diff --git a/deepmd/pt/model/atomic_model/linear_atomic_model.py b/deepmd/pt/model/atomic_model/linear_atomic_model.py
index 41fd49f096..6c620f6f5b 100644
--- a/deepmd/pt/model/atomic_model/linear_atomic_model.py
+++ b/deepmd/pt/model/atomic_model/linear_atomic_model.py
@@ -259,7 +259,7 @@ def forward_atomic(
         """
         nframes, nloc, nnei = nlist.shape
         if (self.do_grad_r() or self.do_grad_c()) and not extended_coord.requires_grad:
-            extended_coord.requires_grad_(True)
+            extended_coord = extended_coord.clone().requires_grad_(True)
         extended_coord = extended_coord.view(nframes, -1, 3)
         sorted_rcuts, sorted_sels = self._sort_rcuts_sels()
         nlists = build_multiple_neighbor_list(
diff --git a/source/tests/pt/model/test_linear_atomic_model.py b/source/tests/pt/model/test_linear_atomic_model.py
index 6e58159ad4..e693e8b90e 100644
--- a/source/tests/pt/model/test_linear_atomic_model.py
+++ b/source/tests/pt/model/test_linear_atomic_model.py
@@ -181,6 +181,22 @@ def test_self_consistency(self) -> None:
             to_numpy_array(ret0["energy"]), ret2["energy"], atol=0.001, rtol=0.001
         )
 
+    def test_forward_atomic_accepts_leaf_view_input(self) -> None:
+        args = [
+            to_torch_tensor(ii) for ii in [self.coord_ext, self.atype_ext, self.nlist]
+        ]
+        coord = args[0]
+        coord_view = coord.view(self.nf, self.nall, 3)
+        coord_view_before = coord_view.detach().clone()
+        self.assertTrue(coord.is_leaf)
+        self.assertTrue(coord_view._is_view())
+        args[0] = coord_view
+        ret = self.md0.forward_atomic(*args)
+
+        self.assertFalse(coord_view.requires_grad)
+        torch.testing.assert_close(coord_view, coord_view_before)
+        self.assertIn("energy", ret)
+
     def test_forward_atomic_preserves_grad_enabled_input(self) -> None:
         args = [
             to_torch_tensor(ii) for ii in [self.coord_ext, self.atype_ext, self.nlist]
diff --git a/source/tests/pt/model/test_pairtab_atomic_model.py b/source/tests/pt/model/test_pairtab_atomic_model.py
index 6c3c325a3e..01927b4967 100644
--- a/source/tests/pt/model/test_pairtab_atomic_model.py
+++ b/source/tests/pt/model/test_pairtab_atomic_model.py
@@ -96,6 +96,18 @@ def test_with_mask(self) -> None:
             result["energy"], expected_result, rtol=0.0001, atol=0.0001
         )
 
+    def test_forward_common_atomic_accepts_leaf_view_input(self) -> None:
+        coord = self.extended_coord.clone()
+        coord_view = coord.view(2, 4, 3)
+        coord_view_before = coord_view.detach().clone()
+        self.assertTrue(coord.is_leaf)
+        self.assertTrue(coord_view._is_view())
+        ret = self.model.forward_atomic(coord_view, self.extended_atype, self.nlist)
+
+        self.assertFalse(coord_view.requires_grad)
+        torch.testing.assert_close(coord_view, coord_view_before)
+        self.assertIn("energy", ret)
+
     def test_forward_common_atomic_preserves_grad_enabled_input(self) -> None:
         coord = self.extended_coord.clone().requires_grad_(True)
         ret = self.model.forward_atomic(coord, self.extended_atype, self.nlist)

From 21cd65aa009908b1ce7bd63f2c909421d0c6445d Mon Sep 17 00:00:00 2001
From: "njzjz-bot (driven by OpenClaw (model:
 custom-chat-jinzhezeng-group/gpt-5.5))[bot]"
 <48687836+njzjz-bot@users.noreply.github.com>
Date: Tue, 2 Jun 2026 04:20:47 +0000
Subject: [PATCH 8/8] style(pt): format DP model grad test

Apply ruff-format's line wrapping to the DP model regression test so pre-commit.ci does not need to autofix the PR branch.

Authored by OpenClaw (model: custom-chat-jinzhezeng-group/gpt-5.5)
---
 source/tests/pt/model/test_dp_model.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/source/tests/pt/model/test_dp_model.py b/source/tests/pt/model/test_dp_model.py
index 63a5b45f34..21335b5bec 100644
--- a/source/tests/pt/model/test_dp_model.py
+++ b/source/tests/pt/model/test_dp_model.py
@@ -151,7 +151,9 @@ def test_forward_lower_accepts_leaf_view_input(self) -> None:
         self.assertIn("virial", ret)
 
         coord_view_grad = coord_ext.view(self.nf, -1, 3).clone().requires_grad_(True)
-        ret = md0.forward_lower(coord_view_grad, atype_ext, nlist, do_atomic_virial=True)
+        ret = md0.forward_lower(
+            coord_view_grad, atype_ext, nlist, do_atomic_virial=True
+        )
 
         self.assertTrue(coord_view_grad.requires_grad)
         self.assertIn("extended_force", ret)