Skip to content

Promote remaining duck-typed attribute probes (add_chg_spin_ebd, geo_compress, reinit_exclude) to base-class interfaces #5897

Closed

Description

@wanghan-iapcm

Follow-up to the capability-method cleanups on the DPA4 graph branch (supports_native_spin/supports_charge_spin, the uses_graph_lower cluster, and has_spin are now declared on the shared base classes with concrete defaults and probed via direct calls).

Call sites use getattr(obj, "attr", <default>) / hasattr duck typing instead of an interface declared once on the owning base class (make_base_descriptor / make_base_fitting / make_base_model / BaseAtomicModel). Each entry below is a candidate for a concrete-default method on its base, so a typo'd or renamed name raises instead of silently degrading to the default.

Inventory

(Updated after a receiver-agnostic re-sweep of both getattr and hasattr over all of deepmd/ outside the frozen deepmd/pt tree — the original issue text understated the scope.)

Charge-spin interface family (largest cluster, ~25 sites)

  • has_chg_spin_ebd — 5 hasattr sites; get_dim_chg_spin — 10 hasattr sites; has_default_chg_spin — 5; get_default_chg_spin — 3; add_chg_spin_ebd attribute — dpmodel/atomic_model/dp_atomic_model.py, pt_expt/infer/deep_eval.py (2 sites). Candidate: declare the whole family on the owning bases (atomic model and/or descriptor) with concrete defaults (False/0/None) and drop every probe.

Compression / kernel-eligibility attributes

  • geo_compress — 12 sites (7 getattr + 5 hasattr) across pt_expt/utils/serialization.py, pt_expt/model/make_model.py, jax/utils/serialization.py, deepmd/kernels/cuda/dpa1/canonical.py. Candidate: base-descriptor capability accessor.
  • _fused_eligible — 4 sites (pt_expt/utils/serialization.py, kernels/triton/dpa1/{edge_conv,se_conv}.py, kernels/cuda/dpa1/canonical.py).
  • tebd_compress, type_embd_data — jax/utils/serialization.py + assorted hasattr sites in dpmodel descriptors.

Fitting interface

  • get_var_name — 9 hasattr sites; get_task_dim — 4; get_intensive — 4. Candidate: concrete defaults on make_base_fitting.
  • reinit_exclude — dpmodel/atomic_model/dp_atomic_model.py via hasattr. Candidate: concrete no-op on make_base_fitting.

Atomic-model attributes

  • pair_excl — 8 getattr(..., None) sites (dpmodel/pt_expt/jax/kernels). Dead-defensive if BaseAtomicModel.__init__ always sets it — verify, then make the direct attribute access the contract.
  • pair_exclude_types — pt_expt/utils/serialization.py.

Dead-defensive probes of already-declared base methods (drop the probe, call directly)

  • has_message_passing — 2 hasattr sites; has_default_fparam — 2 hasattr sites (e.g. jax/jax2tf/tfmodel.py).

Self-probes with defaults (attribute should be declared, not probed)

  • use_loc_mapping — pt_expt/descriptor/repflows.py (getattr(self, ..., False)).
  • set_davg_zero / set_stddev_constant — dpmodel/utils/env_mat_stat.py, pd/model/descriptor/descriptor.py.
  • trainable — probed with INCONSISTENT defaults: True at pt_expt/descriptor/dpa4.py vs False at pt_expt/utils/network.py — an unpinned contract; decide the default once, declare it, and align both.

tf2 optional-feature hooks

  • set_enable_compile, _call_common_lower_formatted — tf2/train/trainer.py (getattr(model, ..., None) then call-if-present). Candidate: declared no-op hooks on the base model.

Explicitly out of scope (legitimate duck typing, not the disease)

  • getattr(torch.ops.deepmd, "...", None) custom-op registry probes (ops genuinely optional at runtime).
  • serialize/deserialize protocol checks in the generic pt_expt wrapper machinery.
  • Probes on external/non-deepmd objects (dtype, shape, dlpack, sess, ...).
  • deepmd/tf (TF1) and deepmd/pt legacy trees (frozen; base declarations still propagate to their classes where the factories are shared).

Notes

  • Defaults must be concrete, not abstract, so descriptors/fittings across all backends need no change (the base factories are shared by dpmodel/pt/pd/tf).
  • Base-class docstrings stay implementer-agnostic (no concrete descriptor names or mechanism names).
  • Each promotion comes with a universal-suite assertion covering both branches (default and override), following DescriptorTestCase.test_capability_contract / ModelTestCase.test_has_spin.
  • The pd has_spin probe trio found in the re-sweep has already been fixed on the DPA4 graph branch (direct calls).

Design caveats (apply per item before promoting)

  1. Guard against base-interface bloat. Every capability method that only a handful of subclasses override widens the shared base interface (the charge-spin family alone adds ~5 methods). Admission criterion: promote only when the probe spans multiple backends/modules; for a single-call-site probe, first consider narrowing the call site instead of widening the base.
  2. Kernel-eligibility attributes may belong on the kernel side, not the descriptor base. _fused_eligible (and to a degree geo_compress) are kernel/backend implementation concerns, not descriptor domain semantics — promoting an underscore-private name into the shared backend-agnostic base leaks kernel concerns into dpmodel. Alternative to evaluate per item: let the kernel layer own the eligibility check ("can I fuse this descriptor?") rather than the descriptor declaring "I am fusable".
  3. pair_excl: verification must precede deletion. Making direct attribute access the contract is valid only after verifying that BaseAtomicModel.__init__ unconditionally sets it on every construction path (including deserialize). Deleting the defensive getattr(..., None) probes first orphans that precondition — enumerate the paths, then drop the probes.

Triage (2026-08-11, applies the caveats above per item)

A. Promote to a base interface

Item Owning base Concrete default Notes
has_chg_spin_ebd / get_dim_chg_spin / get_default_chg_spin make_base_descriptor, forwarded by BaseAtomicModel / make_base_model False / 0 / None Merge has_default_chg_spin into get_default_chg_spin (derive via is not None) — family shrinks 5 → 4 methods. The jax2tf boolean tensor is derived at export time.
get_var_name / get_task_dim / get_intensive make_base_model (get_intensive already concrete at base_atomic_model.py:667) None / per-semantics / False deepmd/infer/deep_eval.py:398 uses not hasattr(get_var_name) to detect non-property models → rewrite as get_var_name() is None value check. [amended during implementation] deepmd/infer/deep_eval.py:398's helper (_get_property_var_name) straddles live models AND loaded artifacts — its hasattr probe is an artifact-boundary check and stays as hasattr, not a value check. The jax2tf/tfmodel.py and tf2/infer/deep_eval.py hasattr(model, "get_var_name"/"get_task_dim"/"get_intensive") sites are likewise artifact-boundary probes (loaded SavedModel/jax artifacts do not always expose these ops) — out of scope, not promotion targets.
geo_compress make_base_descriptor __init__ attribute False Generic consumers: pt_expt/model/make_model.py:783, kernels/cuda/dpa1/canonical.py:37. Pitfall: jax restore walker (jax/utils/serialization.py:78 ff.) uses hasattr as a capability check — semantics invert once the base declares the attribute; rewrite those sites to value checks.
reinit_exclude make_base_fitting concrete no-op Kills the dp_atomic_model.py:120 probe; GeneralFitting.reinit_exclude already exists.
get_pair_exclude_types() (new accessor) BaseAtomicModel concrete return self.pair_exclude_types Serialization/export/jax_md consumers switch to the getter; kernels/cuda/dpa1/canonical.py:42's pair_excl is not None becomes bool(get_pair_exclude_types()) (equivalent: reinit_pair_exclude guarantees empty list ⇔ mask is None).
has_message_passing / has_default_fparam already declared — Drop the 2+2 dead-defensive probes, call directly.
trainable dpmodel network/descriptor __init__ attribute adjudicate first Inconsistent defaults: pt_expt/utils/network.py:102 uses False, the other 9 sites use True. Decide whether False is a typo or deliberate semantics, then declare once.
set_davg_zero / set_stddev_constant descriptor block base (NOT make_base_descriptor) False Probe receivers are blocks (dpmodel/utils/env_mat_stat.py:72,96); declaring on the top-level descriptor base would be interface leakage. [amended during implementation] The "blocks only" premise was false: merge_env_stat's actual contract is Union[Descriptor, DescriptorBlock] — bare se-family descriptors (no block layer) reach merge_env_stat directly via multi-task share_params. Both attributes are therefore also declared, with the same concrete default, on make_base_descriptor's BD base so bare descriptors satisfy the contract too; user-ratified deviation from the original plan premise.

B. Pin the contract (no new method; tighten the invariant)

Item Treatment
pair_excl Verify every construction path (incl. deserialize) goes through BaseAtomicModel.__init__ (which unconditionally calls reinit_pair_exclude) → delete the internal-pipeline getattr(..., None) probes (make_model etc.) in favor of direct attribute access + add a construction-path assertion test. External consumers are taken over by get_pair_exclude_types() above. The nested getattr(model, "atomic_model", None) outer probes are a separate contract — leave them.

C. Narrow fix (family-/call-site-local; do NOT widen the base)

Item Treatment
tebd_compress Declare = False in the dpmodel tebd-family __init__s (dpa1 / se_atten_v2 / se_t_tebd + blocks). The jax walker's hasattr capability check stays valid because only family classes carry the attribute — no rewrite needed.
type_embd_data Same, declare = None in family __init__s (data payload, not a capability flag).
use_loc_mapping dpmodel RepFlows already owns the attribute; pt_expt/descriptor/repflows.py:60 switches to direct access.

D. Leave as-is (kernel-/single-backend-local protocols)

Item Reason
_fused_eligible All 3 sites inside deepmd/kernels/; kernel-side eligibility, underscore-private — caveat 2 verbatim. Optionally consolidate into a kernels-local helper.
set_enable_compile / _call_common_lower_formatted All sites in tf2/train/trainer.py, single-backend optional hooks. Revisit if a second backend needs them.

[amended during implementation] tf2/train/trainer.py's set_enable_compile / _call_common_lower_formatted and _fused_eligible were left as-is per this table, as originally triaged. During the residual audit, the atom_excl getattr siblings (e.g. deepmd/kernels/cuda/dpa1/canonical.py:44) were confirmed to be the same duck-typing disease but are outside this issue's inventory entirely (issue text never listed atom_excl) — candidate follow-up issue, not fixed here.

Cross-cutting

  1. hasattr semantic inversion: the flagged sites in table A cannot be mechanically replaced — each must become a value check.
  2. Every promotion ships the dual-branch (default + override) universal-suite assertion; pair_excl ships the construction-path assertion instead.

Activity

  1. wanghan-iapcm commented on Aug 12, 2026

    @wanghan-iapcm
    CollaboratorAuthor

    All Triage table A/B/C items landed on branch refactor-promote-capability-probes — opened as PR #5967.

    What landed

    • charge-spin family: has_default_chg_spin merged into get_default_chg_spin() (is not None is the predicate); concrete defaults on descriptor base / BaseAtomicModel / make_base_model; all ~25 probes deleted.
    • property trio: get_var_name (→ None) / get_task_dim (→ raise) / get_intensive (→ False) declared on make_base_model; live-model probes deleted.
    • get_geo_compress() method accessor on the descriptor base, with dual-branch test pinning; consumer getattrs deleted.
    • reinit_exclude: concrete default on make_base_fitting — no-op for empty input, raises for non-empty.
    • get_pair_exclude_types() on make_base_atomic_model + pair_excl pinned as a direct-access contract (construction-path assertions; includes SpinModel).
    • set_davg_zero / set_stddev_constant: class-attribute defaults on the DescriptorBlock bases.
    • dead-defensive probes of has_message_passing / has_default_fparam / use_loc_mapping / trainable deleted (fail-fast now).
    • tebd family: type_embd_data / tebd_compress declared in family __init__s; pt_expt register_buffer trap defused (incl. DPA2).

    Triage-table corrections (marked [amended during implementation] in the body)

    1. set_davg_zero/set_stddev_constant are also declared on make_base_descriptor's BD base: merge_env_stat's real contract is Union[Descriptor, DescriptorBlock] — bare se-family descriptors reach it via multi-task share_params, so the original "blocks only" premise was false (user-ratified).
    2. deepmd/infer/deep_eval.py's _get_property_var_name helper straddles live models and loaded artifacts — its hasattr stays as an artifact-boundary check. Same for the jax2tf/tfmodel.py and tf2/infer/deep_eval.py hasattr sites (loaded TF SavedModels): legitimate duck typing, out of scope.
    3. tf2/train/trainer.py's set_enable_compile/_call_common_lower_formatted and _fused_eligible left as-is, per table D.

    Residuals confirmed by the final audit sweep (intentionally kept)

    • atom_excl getattr siblings (e.g. deepmd/kernels/cuda/dpa1/canonical.py:44) — same pattern, but never in this issue's inventory → candidate follow-up issue.
    • jax restore walker (jax/utils/serialization.py) and the trainable tree-walkers — attribute-based family checks, kept by design.
    • Outer getattr(model, "atomic_model", None) shape probes — a separate model-shape contract, unchanged.
    • deepmd/pt / deepmd/tf — frozen trees, untouched.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions