fix(pt_expt): keep the graph-lower routing bool out of traced forwards; CI fixes

Three CI failures from the review-round batch:

1. uses_graph_lower() read the persisted 'graph_lower_disabled' buffer
   via bool(tensor). The predicate runs inside traced forwards (the
   dense-adapter gate in DescrptDPA1.call), so under torch.export the
   read became a data-dependent bool(FakeTensor) guard --
   GuardOnDataDependentSymNode Eq(u0, 1) on every
   test_exportable[excl_types*] parametrization. The routing predicate
   is a plain python bool again; the buffer stays the persistence
   vehicle, re-synced at load time inside _load_from_state_dict where
   the incoming tensor is real. All persistence regressions
   (Trainer restart, state_dict round-trip, pre-knob back-compat) still
   pass.

2. test_dos_graph relied on non-energy models DEFAULT-flipping to the
   graph route, which the energy gate deliberately removed; the parity
   test now opts into the graph route explicitly
   (neighbor_graph_method='dense'), which remains supported for
   non-energy models (eager, output-agnostic).

3. pre-commit: missing ANN annotations on the new
   _load_from_state_dict overrides.

Author: Han Wang

deepmd/pt_expt/descriptor/dpa1.py

@@ -48,29 +48,28 @@ def disable_graph_lower(self) -> None:
         super().disable_graph_lower()
         self.graph_lower_disabled.fill_(True)
 
-    def uses_graph_lower(self) -> bool:
-        """Graph-lower eligibility; the persisted buffer is the source of truth.
-
-        ``load_state_dict`` (Trainer restart) only restores tensor values, so
-        the dpmodel-side bool is re-synced from the buffer here before
-        delegating to the base-class eligibility logic.
-
-        Returns
-        -------
-        bool
-            Whether the descriptor routes through the carry-all graph lower.
-        """
-        if bool(self.graph_lower_disabled):
-            self._graph_lower_disabled = True
-        return super().uses_graph_lower()
-
-    def _load_from_state_dict(self, state_dict, prefix, *args, **kwargs) -> None:
+    def _load_from_state_dict(
+        self,
+        state_dict: dict[str, Any],
+        prefix: str,
+        *args: Any,
+        **kwargs: Any,
+    ) -> None:
         # Back-compat: checkpoints written before the knob was persisted lack
         # the buffer; default to the fresh module's value (graph enabled)
         # instead of failing the strict load.
         key = prefix + "graph_lower_disabled"
         if key not in state_dict:
             state_dict[key] = self.graph_lower_disabled.detach().clone()
+        else:
+            # Re-sync the dpmodel-side routing bool from the RESTORED value
+            # here, at load time, where the incoming tensor is real.  The
+            # routing predicate itself must stay a plain python bool:
+            # ``uses_graph_lower()`` runs inside traced forwards (the dense
+            # adapter gate), and reading the buffer there would emit a
+            # data-dependent ``bool(FakeTensor)`` guard that breaks
+            # torch.export (GuardOnDataDependentSymNode Eq(u0, 1)).
+            self._graph_lower_disabled = bool(state_dict[key])
         super()._load_from_state_dict(state_dict, prefix, *args, **kwargs)
 
     def share_params(

deepmd/pt_expt/descriptor/dpa2.py

@@ -52,29 +52,28 @@ def disable_graph_lower(self) -> None:
         super().disable_graph_lower()
         self.graph_lower_disabled.fill_(True)
 
-    def uses_graph_lower(self) -> bool:
-        """Graph-lower eligibility; the persisted buffer is the source of truth.
-
-        ``load_state_dict`` (Trainer restart) only restores tensor values, so
-        the dpmodel-side bool is re-synced from the buffer here before
-        delegating to the base-class eligibility logic.
-
-        Returns
-        -------
-        bool
-            Whether the descriptor routes through the carry-all graph lower.
-        """
-        if bool(self.graph_lower_disabled):
-            self._graph_lower_disabled = True
-        return super().uses_graph_lower()
-
-    def _load_from_state_dict(self, state_dict, prefix, *args, **kwargs) -> None:
+    def _load_from_state_dict(
+        self,
+        state_dict: dict[str, Any],
+        prefix: str,
+        *args: Any,
+        **kwargs: Any,
+    ) -> None:
         # Back-compat: checkpoints written before the knob was persisted lack
         # the buffer; default to the fresh module's value (graph enabled)
         # instead of failing the strict load.
         key = prefix + "graph_lower_disabled"
         if key not in state_dict:
             state_dict[key] = self.graph_lower_disabled.detach().clone()
+        else:
+            # Re-sync the dpmodel-side routing bool from the RESTORED value
+            # here, at load time, where the incoming tensor is real.  The
+            # routing predicate itself must stay a plain python bool:
+            # ``uses_graph_lower()`` runs inside traced forwards (the dense
+            # adapter gate), and reading the buffer there would emit a
+            # data-dependent ``bool(FakeTensor)`` guard that breaks
+            # torch.export (GuardOnDataDependentSymNode Eq(u0, 1)).
+            self._graph_lower_disabled = bool(state_dict[key])
         super()._load_from_state_dict(state_dict, prefix, *args, **kwargs)
 
     def share_params(

source/tests/pt_expt/model/test_dos_graph.py

@@ -5,7 +5,8 @@
 routes into the graph path by default.  Before the general output transform this
 KeyError'd on ``"energy"``; now every fitting (dos/dipole/polar/property/...)
 flows through :func:`fit_output_to_model_output_graph` with no change on the
-fitting side.  Each model's graph forward (default ``neighbor_graph_method``)
+fitting side.  Each model's graph forward (EXPLICIT ``neighbor_graph_method``
+opt-in; non-energy models default to dense since the energy gate)
 must match the dense path (``neighbor_graph_method="legacy"``) on every shared
 key (carry-all graph at non-binding ``sel`` reproduces the dense neighbor set).
 """
@@ -129,8 +130,14 @@ def test_dos_repro(self) -> None:
         [_make_dos, _make_dipole, _make_polar, _make_property],
     )  # one builder per fitting kind
     def test_graph_matches_dense(self, make_model) -> None:
-        """Graph (default) output matches the dense (``legacy``) path on every
+        """EXPLICIT graph opt-in matches the dense (``legacy``) path on every
         shared key, including derivatives for r/c-differentiable fittings.
+
+        Non-energy models no longer DEFAULT-flip to the graph route (the
+        compiled-training trace is energy-specific; see the
+        ``_resolve_graph_method`` energy gate), but the eager graph lower
+        stays output-agnostic and available via an explicit
+        ``neighbor_graph_method`` -- which is what this parity test pins.
         """
         tol = (
             {"rtol": 1e-11, "atol": 1e-11}
@@ -139,7 +146,13 @@ def test_graph_matches_dense(self, make_model) -> None:
         )
         ds = _make_descriptor()
         m = make_model(ds)
-        graph = m.call_common(self.coord, self.atype, None, do_atomic_virial=True)
+        graph = m.call_common(
+            self.coord,
+            self.atype,
+            None,
+            do_atomic_virial=True,
+            neighbor_graph_method="dense",
+        )
         # the dense path differentiates w.r.t. coord -> needs a coord leaf.
         dense = m.call_common(
             self.coord.detach().requires_grad_(True),