fix(pt-expt): preserve lower semantics in backend conversion

deepmd/entrypoints/convert_backend.py

@@ -30,6 +30,14 @@ def convert_backend(
         If True, export .pt2/.pte models with per-atom virial correction.
         This adds ~2.5x inference cost.  Default False.  Silently ignored
         (with a warning) for backends that don't support the flag.
+
+    Notes
+    -----
+    Backend conversion preserves an explicit ``lower_input_kind`` reported by
+    the source serializer. Sources without this metadata retain the target's
+    automatic lower selection for backward compatibility. A target backend
+    that cannot represent an explicit non-dense lower is rejected rather than
+    silently changing the model function.
     """
     inp_backend: Backend = Backend.detect_backend_by_model(INPUT)()
     out_backend: Backend = Backend.detect_backend_by_model(OUTPUT)()
@@ -40,8 +48,18 @@ def convert_backend(
 
     sig = inspect.signature(out_hook)
     hook_kwargs: dict[str, Any] = {}
+    lower_input_kind = data.get("lower_input_kind")
     if "lower_kind" in sig.parameters:
-        hook_kwargs["lower_kind"] = "auto"
+        hook_kwargs["lower_kind"] = (
+            lower_input_kind if lower_input_kind is not None else "auto"
+        )
+    elif lower_input_kind not in (None, "nlist"):
+        raise ValueError(
+            f"Cannot preserve lower_input_kind {lower_input_kind!r} when "
+            f"converting to output backend {out_backend.name!r}: its "
+            "deserializer does not accept a lower_kind. Retrain or freeze the "
+            "model with that backend instead of converting this artifact."
+        )
     if "do_atomic_virial" in sig.parameters:
         hook_kwargs["do_atomic_virial"] = atomic_virial
     elif atomic_virial:

deepmd/jax/utils/serialization.py

@@ -427,6 +427,7 @@ def restore_model(model_params: dict, model_state: dict) -> BaseModel:
             "jax_version": jax.__version__,
             "model": model_dict,
             "model_def_script": model_def_script,
+            "lower_input_kind": "nlist",
             "@variables": {},
         }
         if min_nbor_dist is not None:
@@ -436,6 +437,7 @@ def restore_model(model_params: dict, model_state: dict) -> BaseModel:
         data = load_dp_model(model_file)
         data.pop("constants")
         data["@variables"].pop("stablehlo")
+        data["lower_input_kind"] = "nlist"
         return data
     elif model_file.endswith(".savedmodel"):
         raise ValueError(

deepmd/pt/model/model/model.py

@@ -26,6 +26,17 @@ def __init__(self, *args: Any, **kwargs: Any) -> None:
             "min_nbor_dist", torch.tensor(-1.0, dtype=torch.float64, device=env.DEVICE)
         )
 
+    def export_lower_input_kind(self) -> str:
+        """Return the lower-input ABI that preserves this model's semantics.
+
+        Returns
+        -------
+        str
+            ``"nlist"`` for the standard PyTorch model contract. Models with
+            a graph-native deployment ABI override this method.
+        """
+        return "nlist"
+
     def compute_or_load_stat(
         self,
         sampled_func: Any,

deepmd/pt/model/model/spin_model.py

@@ -454,6 +454,17 @@ def has_spin(self) -> bool:
         """Returns whether it has spin input and output."""
         return True
 
+    def export_lower_input_kind(self) -> str:
+        """Return the dense ABI used by the virtual-atom spin model.
+
+        Returns
+        -------
+        str
+            ``"nlist"``, because virtual atoms are expanded inside the
+            bounded neighbor-list contract.
+        """
+        return "nlist"
+
     @torch.jit.export
     def has_message_passing(self) -> bool:
         """Returns whether the model has message passing."""

deepmd/pt/utils/serialization.py

@@ -53,6 +53,7 @@ def serialize_from_file(model_file: str) -> dict:
         "pt_version": str(torch.__version__),
         "model": model_dict,
         "model_def_script": model_def_script,
+        "lower_input_kind": model.export_lower_input_kind(),
         "@variables": {},
     }
     if model.get_min_nbor_dist() is not None:

deepmd/pt_expt/utils/serialization.py

@@ -52,6 +52,18 @@
 # ---------------------------------------------------------------------------
 PT2_EXTRA_PREFIX = "model/extra/"
 
+# Backend conversion supplies the source artifact's lower ABI. Each accepted
+# source kind maps to the concrete schema emitted by pt_expt. PT SeZM's
+# ``edge_vec`` lower and pt_expt's ``graph`` lower carry the same directed-edge
+# model semantics, while the target materializes its native NeighborGraph ABI.
+_LOWER_INPUT_KIND_TARGETS = {
+    "nlist": "nlist",
+    "graph": "graph",
+    "dpa1_canonical": "dpa1_canonical",
+    "dpa4c_canonical": "dpa4c_canonical",
+    "edge_vec": "graph",
+}
+
 
 def _strip_shape_assertions(graph_module: torch.nn.Module) -> None:
     """Neutralise deferred shape-guard assertion nodes in an exported graph.
@@ -1242,7 +1254,8 @@ def serialize_from_file(model_file: str) -> dict:
     dict
         The serialized model data.  If the archive contains
         ``model_def_script.json`` (training config), it is included
-        under the ``"model_def_script"`` key.
+        under the ``"model_def_script"`` key. ``lower_input_kind`` records
+        the concrete lower ABI from the artifact metadata.
     """
     if model_file.endswith(".pt2"):
         return _serialize_from_file_pt2(model_file)
@@ -1252,10 +1265,20 @@ def serialize_from_file(model_file: str) -> dict:
 
 def _serialize_from_file_pte(model_file: str) -> dict:
     """Serialize a .pte model file to a dictionary."""
-    extra_files = {"model.json": "", "model_def_script.json": ""}
+    extra_files = {
+        "model.json": "",
+        "model_def_script.json": "",
+        "metadata.json": "",
+    }
     torch.export.load(model_file, extra_files=extra_files)
     model_dict = json.loads(extra_files["model.json"])
     model_dict = _json_to_numpy(model_dict)
+    metadata = (
+        json.loads(extra_files["metadata.json"]) if extra_files["metadata.json"] else {}
+    )
+    model_dict["lower_input_kind"] = metadata.get(
+        "lower_input_kind", model_dict.get("lower_input_kind", "nlist")
+    )
     if extra_files["model_def_script.json"]:
         model_dict["model_def_script"] = json.loads(
             extra_files["model_def_script.json"]
@@ -1273,6 +1296,7 @@ def _serialize_from_file_pt2(model_file: str) -> dict:
 
     model_json_entry = PT2_EXTRA_PREFIX + "model.json"
     model_def_script_entry = PT2_EXTRA_PREFIX + "model_def_script.json"
+    metadata_entry = PT2_EXTRA_PREFIX + "metadata.json"
     with zipfile.ZipFile(model_file, "r") as zf:
         names = zf.namelist()
         if model_json_entry not in names:
@@ -1283,8 +1307,15 @@ def _serialize_from_file_pt2(model_file: str) -> dict:
         model_def_script_json = ""
         if model_def_script_entry in names:
             model_def_script_json = zf.read(model_def_script_entry).decode("utf-8")
+        metadata_json = ""
+        if metadata_entry in names:
+            metadata_json = zf.read(metadata_entry).decode("utf-8")
     model_dict = json.loads(model_json)
     model_dict = _json_to_numpy(model_dict)
+    metadata = json.loads(metadata_json) if metadata_json else {}
+    model_dict["lower_input_kind"] = metadata.get(
+        "lower_input_kind", model_dict.get("lower_input_kind", "nlist")
+    )
     if model_def_script_json:
         model_dict["model_def_script"] = json.loads(model_def_script_json)
     return model_dict
@@ -1355,6 +1386,17 @@ def _resolve_lower_kind(model_file: str, data: dict, lower_kind: str) -> str:
     return "nlist"
 
 
+def _resolve_target_lower_kind(model_file: str, data: dict, lower_kind: str) -> str:
+    """Resolve a source lower ABI to a concrete pt_expt export schema."""
+    source_lower_kind = _resolve_lower_kind(model_file, data, lower_kind)
+    if source_lower_kind not in _LOWER_INPUT_KIND_TARGETS:
+        raise ValueError(
+            f"Unsupported lower_kind {source_lower_kind!r}; expected one of "
+            f"{sorted(_LOWER_INPUT_KIND_TARGETS)}."
+        )
+    return _LOWER_INPUT_KIND_TARGETS[source_lower_kind]
+
+
 def deserialize_to_file(
     model_file: str,
     data: dict,
@@ -1393,14 +1435,17 @@ def deserialize_to_file(
         (``atype``/``n_node``/``edge_index``/``edge_vec``/``edge_mask`` and
         the destination/source CSR views) with a DYNAMIC edge axis ``E``
         (``Dim("nedge", min=2)``), so the artifact accepts any system size.
-        ``"auto"`` (used by ``convert-backend``) resolves to ``"graph"`` for an
-        exportable graph-lower ``.pt2`` and ``"nlist"`` otherwise (see
-        :func:`_resolve_lower_kind`). A graph lower always preserves the fused
-        inference operators (``DP_CUDA_INFER >= 2``) and the per-atom virial.
+        ``"auto"`` resolves to ``"graph"`` for an exportable graph-lower
+        ``.pt2`` and ``"nlist"`` otherwise (see :func:`_resolve_lower_kind`).
+        Backend conversion passes the source artifact's concrete lower kind;
+        compatible source ABIs are mapped to the target's native schema while
+        preserving their execution semantics. A graph lower always preserves
+        the fused inference operators (``DP_CUDA_INFER >= 2``) and the
+        per-atom virial.
         The selected schema is recorded as ``lower_input_kind`` in
         ``metadata.json``.
     """
-    lower_kind = _resolve_lower_kind(model_file, data, lower_kind)
+    lower_kind = _resolve_target_lower_kind(model_file, data, lower_kind)
     if data["model"].get("type") == "native_spin" and lower_kind not in (
         "graph",
         "dpa4c_canonical",

deepmd/tf/utils/serialization.py

@@ -50,6 +50,7 @@ def serialize_from_file(model_file: str) -> dict:
         "tf_version": tf.__version__,
         "model": model_dict,
         "model_def_script": jdata["model"],
+        "lower_input_kind": "nlist",
     }
     # neighbor stat information
     try:

deepmd/tf2/utils/serialization.py

@@ -611,6 +611,7 @@ def serialize_from_file(model_file: str) -> dict:
         "backend": "TensorFlow2",
         "model": model_payload,
         "model_def_script": model_def_script,
+        "lower_input_kind": "nlist",
         "shared_links": state.get("shared_links"),
         "@variables": {
             "current_step": int(state.get("current_step", 0)),

doc/backend.md

@@ -154,3 +154,12 @@ then selects its `edge_vec`, dense `nlist`, NeighborGraph, or compact
 ## Convert model files between backends
 
 If a model is supported by two backends, one can use [`dp convert-backend`](./cli.rst) to convert the model file between these two backends.
+
+Backend conversion preserves the concrete `lower_input_kind` reported by the
+source serializer. Dense TensorFlow, TensorFlow 2, JAX, and standard PyTorch
+models therefore remain dense `nlist` models when converted to `.pt2`; model
+families with a graph-native deployment ABI report their corresponding kind.
+Compiled `.pt2` and `.pte` artifacts retain the kind recorded in their metadata.
+A conversion is rejected when the target backend cannot represent an explicit
+source kind. Legacy model files without lower metadata retain target-specific
+automatic selection for backward compatibility.

source/tests/consistent/io/test_io.py

@@ -275,6 +275,7 @@ def setUp(self) -> None:
             "model": model.serialize(),
             "backend": "test",
             "model_def_script": model_def_script,
+            "lower_input_kind": "nlist",
         }
 
     def tearDown(self) -> None:
@@ -319,6 +320,7 @@ def setUp(self) -> None:
             "model": model.serialize(),
             "backend": "test",
             "model_def_script": model_def_script,
+            "lower_input_kind": "nlist",
         }
 
     def tearDown(self) -> None:
@@ -375,6 +377,7 @@ def setUp(self) -> None:
             "model": model.serialize(),
             "backend": "test",
             "model_def_script": model_def_script,
+            "lower_input_kind": "nlist",
         }
 
     def tearDown(self) -> None:
@@ -422,6 +425,7 @@ def setUp(self) -> None:
             "model": model.serialize(),
             "backend": "test",
             "model_def_script": model_def_script,
+            "lower_input_kind": "nlist",
         }
 
     def tearDown(self) -> None:

source/tests/jax/test_hlo.py

@@ -1,9 +1,21 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
 """Regression tests for metadata exposed by serialized JAX HLO models."""
 
+from types import (
+    SimpleNamespace,
+)
+
+import pytest
+from typing_extensions import (
+    Self,
+)
+
 from deepmd.jax.model.hlo import (
     HLO,
 )
+from deepmd.jax.utils import (
+    serialization,
+)
 
 
 def test_hlo_get_nnei_uses_stored_selection() -> None:
@@ -17,3 +29,47 @@ def test_hlo_get_nnei_uses_stored_selection() -> None:
     model.sel = [6, 12, 1]
 
     assert model.get_nnei() == sum(model.sel)
+
+
+def test_hlo_serialization_declares_dense_lower(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    """A JAX HLO artifact exposes its dense source execution semantics."""
+    stored_data = {
+        "model": {},
+        "constants": {},
+        "@variables": {"stablehlo": b"module"},
+    }
+    monkeypatch.setattr(serialization, "load_dp_model", lambda _path: stored_data)
+
+    data = serialization.serialize_from_file("model.hlo")
+
+    assert data["lower_input_kind"] == "nlist"
+
+
+def test_checkpoint_serialization_declares_dense_lower(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    """A JAX training checkpoint exposes its dense source semantics."""
+
+    class Checkpointer:
+        def __init__(self, _handler: object) -> None:
+            pass
+
+        def __enter__(self) -> Self:
+            return self
+
+        def __exit__(self, *_args: object) -> None:
+            pass
+
+        def restore(self, *_args: object, **_kwargs: object) -> SimpleNamespace:
+            return SimpleNamespace(
+                state={},
+                model_def_script={"model_dict": {}},
+            )
+
+    monkeypatch.setattr(serialization.ocp, "Checkpointer", Checkpointer)
+
+    data = serialization.serialize_from_file("model.jax")
+
+    assert data["lower_input_kind"] == "nlist"

source/tests/pt/model/test_ener_spin_model.py

@@ -111,6 +111,10 @@ def test_output_shape(
         torch.testing.assert_close(result["force"].shape, [nframes, nloc, 3])
         torch.testing.assert_close(result["force_mag"].shape, [nframes, nloc, 3])
 
+    def test_export_lower_input_kind(self) -> None:
+        """Virtual-atom spin models retain the dense export ABI."""
+        self.assertEqual(self.model.export_lower_input_kind(), "nlist")
+
     def test_input_output_process(self) -> None:
         nframes, nloc = self.coord.shape[:2]
         self.real_ntypes = self.model.spin.get_ntypes_real()