fix: three canonical-bridging P1s from review round 2

1. Sugar expansion keeps trainer-owned top-level `lora` at the
   composition level instead of forwarding it to the learned child
   (which the pt bridge builder rejects), restoring the concise
   dpa4+bridging_method+lora form. The key routing is now a four-way
   table (composition / consumed / trainer / learned-child) whose
   coverage of the standard+dpa4 argcheck schemas is pinned by
   test_routing_covers_the_argcheck_schema: a new argcheck model key
   fails the test until it gets an explicit routing decision, instead
   of silently landing on the child. The guard immediately surfaced
   use_compile and enable_tf32 as previously-undecided keys (both
   routed to the learned child, preserving behavior).

2. pt LinearEnergyModel.update_sel: the shared-config reconstruction
   loop now skips the `inner_potential` child too. Normalization
   always inserts `shared_dict: {}`, so the default CLI path entered
   the reconstruction loop and dereferenced the analytical child's
   missing descriptor (KeyError). New test runs on a NORMALIZED config.

3. The shared dpmodel/pt_expt linear builder rejects bridged
   compositions with a third child (e.g. pairtab): no common execution
   route exists (pairtab is dense-only, the bridged pair graph-only),
   matching the pt builder's exact-two-child constraint and message.

Author: Han Wang

deepmd/dpmodel/model/model_factory.py

@@ -214,11 +214,14 @@ def get_linear_atomic_model(
                 "A linear_ener composition supports at most one "
                 "`inner_potential` sub-model."
             )
-        if len(learned_indices) != 1:
+        if len(learned_indices) != 1 or len(children) != 2:
+            # A third child (e.g. pairtab) has no common execution route
+            # with the graph-only bridged pair; reject at construction
+            # like the pt builder does.
             raise ValueError(
                 "An `inner_potential` sub-model bridges exactly one learned "
-                f"sibling, but got {len(learned_indices)} sub-models with a "
-                "descriptor."
+                "sibling: expected a linear_ener composition over "
+                "[learned, inner_potential]."
             )
         if str(data.get("weights", "mean")) != "sum":
             raise ValueError(

deepmd/pt/model/model/dp_linear_model.py

@@ -370,6 +370,9 @@ def get_shared_key(shared_ref: str) -> str:
         if "type_map" not in ret_jdata:
             ret_jdata["type_map"] = deepcopy(type_map)
         for idx, original_sub_model in enumerate(original_models):
+            if original_sub_model.get("type") == "inner_potential":
+                # analytical child: no descriptor to write back
+                continue
             if "tab_file" in original_sub_model:
                 continue
             updated_sub_model = local_jdata_cpy["models"][idx]

deepmd/utils/bridging.py

@@ -66,17 +66,57 @@ def _is_dpa4_family(sub: dict) -> bool:
     return any(_is_dpa4_family(sub) for sub in children)
 
 
-# Top-level keys that belong to the composition, not to the learned child.
+# Routing of the concise-form top-level keys during sugar expansion. Every
+# key the `standard`/`dpa4` argcheck schemas declare must appear in exactly
+# one tuple below: a schema-coverage test derives the key universe from
+# `deepmd.utils.argcheck` and fails when a new key is left unrouted, so
+# adding a model key forces an explicit routing decision here.
+
+# Keys that belong to the composition, not to the learned child.
 _COMPOSITION_KEYS = (
     "type",
     "type_map",
     "spin",
     "atom_exclude_types",
     "pair_exclude_types",
+)
+# Keys consumed by the expansion itself; they appear in neither the
+# composition nor the learned child.
+_CONSUMED_KEYS = (
     "bridging_method",
     "bridging_r_inner",
     "bridging_r_outer",
 )
+# Training-owned keys: the trainer reads them from the top level of the
+# model section, so they stay at the composition level and must never be
+# forwarded to a sub-model.
+_TRAINER_KEYS = ("lora",)
+# Keys that configure the learned model and are forwarded to the learned
+# child. This tuple is not consulted at expansion time (the child receives
+# every key not routed above); it exists so the schema-coverage test can
+# assert that every argcheck key has an explicit routing decision.
+_LEARNED_CHILD_KEYS = (
+    "descriptor",
+    "fitting_net",
+    "model_branch_alias",
+    "info",
+    "use_compile",
+    "enable_tf32",
+    "data_stat_nbatch",
+    "data_stat_protect",
+    "data_bias_nsample",
+    "use_srtab",
+    "smin_alpha",
+    "sw_rmin",
+    "sw_rmax",
+    "preset_out_bias",
+    "srtab_add_bias",
+    "type_embedding",
+    "modifier",
+    "compress",
+    "finetune_head",
+)
+_NON_CHILD_KEYS = _COMPOSITION_KEYS + _CONSUMED_KEYS + _TRAINER_KEYS
 
 
 def expand_bridging_method(data: dict) -> dict:
@@ -87,8 +127,9 @@ def expand_bridging_method(data: dict) -> dict:
     deep-copied and rewritten to the canonical composition form: a
     ``linear_ener`` model with ``weights: "sum"`` over the learned
     sub-model and an ``inner_potential`` sub-model. The exclusion lists
-    move to the composition level; a top-level ``spin`` section stays at
-    the top level; every other key stays on the learned child.
+    move to the composition level; a top-level ``spin`` section and the
+    training-owned keys (``lora``) stay at the top level; every other key
+    stays on the learned child.
 
     For backward compatibility with the legacy pt ``type: "dpa4"``
     builder, ``descriptor.exclude_types`` is promoted to the composition's
@@ -142,9 +183,7 @@ def expand_bridging_method(data: dict) -> dict:
     else:
         pair_exclude_types = descriptor_exclude_types
 
-    learned = {
-        key: value for key, value in data.items() if key not in _COMPOSITION_KEYS
-    }
+    learned = {key: value for key, value in data.items() if key not in _NON_CHILD_KEYS}
     learned["type"] = model_type
     learned["type_map"] = copy.deepcopy(data["type_map"])
     canonical = {
@@ -165,4 +204,7 @@ def expand_bridging_method(data: dict) -> dict:
     }
     if "spin" in data:
         canonical["spin"] = data["spin"]
+    for key in _TRAINER_KEYS:
+        if key in data:
+            canonical[key] = data[key]
     return canonical

source/tests/common/dpmodel/test_zbl_bridging.py

@@ -76,6 +76,29 @@ def test_builder_composes_linear_model():
     assert float(dp_child.descriptor.inner_clamp.r_inner) == 0.8
 
 
+def test_third_child_without_common_route_raises():
+    """[learned, inner_potential, pairtab] has no common execution route
+    (pairtab is dense-only, the bridged pair is graph-only): the builder
+    must reject it at construction like the pt backend does.
+    """
+    cfg = {
+        "type": "linear_ener",
+        "weights": "sum",
+        "type_map": ["Ni", "O"],
+        "models": [
+            {
+                "type": "dpa4",
+                "descriptor": copy.deepcopy(ZBL_CONFIG["descriptor"]),
+                "fitting_net": copy.deepcopy(ZBL_CONFIG["fitting_net"]),
+            },
+            {"type": "inner_potential", "mode": "ZBL"},
+            {"type": "pairtab", "tab_file": "unused.txt", "rcut": 4.0, "sel": 8},
+        ],
+    }
+    with pytest.raises(ValueError, match="exactly one learned"):
+        get_model(cfg)
+
+
 def test_zbl_child_equals_composition_minus_learned():
     """Composition energy == learned child + analytical child (exact sum)."""
     model = get_model(copy.deepcopy(ZBL_CONFIG))

source/tests/common/test_bridging.py

@@ -103,6 +103,60 @@ def test_other_model_keys_stay_on_learned_child() -> None:
     assert "preset_out_bias" not in out
 
 
+def test_lora_stays_top_level() -> None:
+    """`lora` is training-owned: the pt trainer reads it from the top
+    level of the model section, so the expansion must keep it there and
+    never forward it to the learned child (which the pt bridge builder
+    rejects).
+    """
+    data = _flag_config()
+    data["lora"] = {"rank": 2, "alpha": None}
+    out = expand_bridging_method(data)
+    assert out["lora"] == {"rank": 2, "alpha": None}
+    assert "lora" not in out["models"][0]
+
+
+def test_routing_covers_the_argcheck_schema() -> None:
+    """Every key the `standard`/`dpa4` argcheck schemas declare must have
+    an explicit routing decision in the expansion. Adding a model key to
+    argcheck without deciding its routing fails here instead of silently
+    landing on the learned child (how top-level `lora` once broke).
+    """
+    from deepmd.utils.argcheck import (
+        model_args,
+        sezm_model_args,
+        standard_model_args,
+    )
+    from deepmd.utils.bridging import (
+        _COMPOSITION_KEYS,
+        _CONSUMED_KEYS,
+        _LEARNED_CHILD_KEYS,
+        _TRAINER_KEYS,
+    )
+
+    schema_keys = {"type"}
+    schema_keys |= set(model_args(exclude_hybrid=True).sub_fields)
+    schema_keys |= set(standard_model_args().sub_fields)
+    schema_keys |= set(sezm_model_args().sub_fields)
+
+    routing = [
+        set(_COMPOSITION_KEYS),
+        set(_CONSUMED_KEYS),
+        set(_TRAINER_KEYS),
+        set(_LEARNED_CHILD_KEYS),
+    ]
+    routed = set().union(*routing)
+    assert sum(len(s) for s in routing) == len(routed), (
+        "a key is routed to more than one destination"
+    )
+    assert schema_keys == routed, (
+        f"unrouted argcheck keys: {sorted(schema_keys - routed)}; "
+        f"routed keys absent from the schema: {sorted(routed - schema_keys)}. "
+        "Decide the routing in deepmd.utils.bridging and update the "
+        "corresponding tuple."
+    )
+
+
 def test_exclusions_move_to_composition_level() -> None:
     data = _flag_config()
     data["pair_exclude_types"] = [[0, 1]]

source/tests/pt/model/test_get_model_bridging.py

@@ -252,6 +252,57 @@ def test_is_sezm_checkpoint_recognizes_canonical_params(tmp_path) -> None:
     assert is_sezm_checkpoint(ckpt2)
 
 
+def test_sugar_with_top_level_lora_builds() -> None:
+    """The concise dpa4+bridging form with trainer-owned top-level `lora`
+    must keep building: the expansion routes `lora` to the composition
+    level, so the bridge builder never sees it on the learned child.
+    """
+    from deepmd.pt.train.training import (
+        get_model_for_wrapper,
+    )
+
+    cfg = _sugar_config()
+    cfg["lora"] = {"rank": 2, "alpha": None}
+    model = get_model_for_wrapper(copy.deepcopy(cfg))
+    assert isinstance(model, SeZMModel)
+    # The trainer injects the adapters later by reading the top level of
+    # its own (unexpanded) config; expansion must not have mutated it.
+    assert cfg["lora"] == {"rank": 2, "alpha": None}
+
+
+def test_update_sel_normalized_config_skips_inner_potential_child(
+    monkeypatch,
+) -> None:
+    """The default CLI path hands `update_sel` a NORMALIZED config, where
+    argcheck always inserts `shared_dict: {}`. Both the update loop and
+    the shared-config reconstruction loop must skip the analytical child.
+    """
+    from deepmd.pt.model.model import (
+        LinearEnergyModel,
+    )
+    from deepmd.pt.model.model.dp_model import (
+        DPModelCommon,
+    )
+    from deepmd.utils.argcheck import (
+        model_args,
+    )
+
+    seen = []
+
+    def _fake_update_sel(train_data, type_map, sub):
+        seen.append(copy.deepcopy(sub))
+        return sub, 0.9
+
+    monkeypatch.setattr(DPModelCommon, "update_sel", staticmethod(_fake_update_sel))
+    cfg = model_args().normalize_value(_canonical_config(), trim_pattern="_*")
+    assert cfg["shared_dict"] == {}  # inserted by normalization
+    updated, min_dist = LinearEnergyModel.update_sel(None, cfg["type_map"], cfg)
+    assert min_dist == 0.9
+    assert len(seen) == 1  # only the learned child
+    assert "descriptor" in seen[0]
+    assert updated["models"][1]["type"] == "inner_potential"
+
+
 def test_update_sel_skips_inner_potential_child(monkeypatch) -> None:
     """Neighbor-stat selection must skip the analytical child instead of
     crashing on its missing descriptor.

source/tests/pt_expt/model/test_get_model_bridging.py

@@ -194,6 +194,19 @@ def _canonical_config() -> dict:
     }
 
 
+def test_third_child_without_common_route_raises() -> None:
+    """[learned, inner_potential, pairtab] has no common execution route
+    (pairtab is dense-only, the bridged pair is graph-only): the shared
+    builder must reject it at construction like the pt backend does.
+    """
+    cfg = _canonical_config()
+    cfg["models"].append(
+        {"type": "pairtab", "tab_file": "unused.txt", "rcut": 4.0, "sel": 8}
+    )
+    with pytest.raises(ValueError, match="exactly one learned"):
+        get_model(cfg)
+
+
 def test_canonical_composition_builds() -> None:
     """The canonical spelling composes [learned, InnerPotential] with the
     clamp radii derived onto the learned child's descriptor.