diff --git a/deepmd/dpmodel/atomic_model/base_atomic_model.py b/deepmd/dpmodel/atomic_model/base_atomic_model.py index 98e6e0cd3a..12be2aafe2 100644 --- a/deepmd/dpmodel/atomic_model/base_atomic_model.py +++ b/deepmd/dpmodel/atomic_model/base_atomic_model.py @@ -250,10 +250,6 @@ def get_dim_chg_spin(self) -> int: """Get the dimension of charge_spin input.""" return 0 - def has_default_chg_spin(self) -> bool: - """Check if the model has default charge_spin values.""" - return False - def get_default_chg_spin(self) -> list[float] | None: """Get the default charge_spin values.""" return None diff --git a/deepmd/dpmodel/atomic_model/dp_atomic_model.py b/deepmd/dpmodel/atomic_model/dp_atomic_model.py index c5cf0b157c..0406eb4b91 100644 --- a/deepmd/dpmodel/atomic_model/dp_atomic_model.py +++ b/deepmd/dpmodel/atomic_model/dp_atomic_model.py @@ -117,12 +117,9 @@ def __init__( super().__init__(type_map, **kwargs) self.descriptor = descriptor self.fitting_net = fitting - if hasattr(self.fitting_net, "reinit_exclude"): - self.fitting_net.reinit_exclude(self.atom_exclude_types) + self.fitting_net.reinit_exclude(self.atom_exclude_types) self.type_map = type_map - self.add_chg_spin_ebd: bool = getattr( - self.descriptor, "add_chg_spin_ebd", False - ) + self.add_chg_spin_ebd: bool = self.descriptor.get_dim_chg_spin() > 0 # Structural capability: only descriptors with a native spin # conditioning mechanism (currently DPA4) accept a ``spin`` kwarg on # ``call_graph`` at all -- unlike ``charge_spin``, which every @@ -151,15 +148,9 @@ def get_dim_chg_spin(self) -> int: return self.descriptor.get_dim_chg_spin() return 0 - def has_default_chg_spin(self) -> bool: - """Check if the model has default charge_spin values.""" - if self.add_chg_spin_ebd: - return self.descriptor.has_default_chg_spin() - return False - def get_default_chg_spin(self) -> list[float] | None: """Get the default charge_spin values.""" - if self.add_chg_spin_ebd and self.descriptor.has_default_chg_spin(): + if self.add_chg_spin_ebd: return self.descriptor.get_default_chg_spin() return None diff --git a/deepmd/dpmodel/atomic_model/linear_atomic_model.py b/deepmd/dpmodel/atomic_model/linear_atomic_model.py index 4999bc15a7..bb42ec37ae 100644 --- a/deepmd/dpmodel/atomic_model/linear_atomic_model.py +++ b/deepmd/dpmodel/atomic_model/linear_atomic_model.py @@ -723,19 +723,11 @@ def _fparam_consumers(self) -> list: """Children that actually consume ``fparam``.""" return [m for m in self.models if m.get_dim_fparam() > 0] - def has_default_chg_spin(self) -> bool: - """Whether every active child shares one default charge/spin.""" - return self._agreed_default( - self._chg_spin_consumers(), - lambda m: m.has_default_chg_spin(), - lambda m: m.get_default_chg_spin(), - )[0] - def get_default_chg_spin(self) -> "Array | None": """The shared default charge/spin conditions, if the children agree.""" return self._agreed_default( self._chg_spin_consumers(), - lambda m: m.has_default_chg_spin(), + lambda m: m.get_default_chg_spin() is not None, lambda m: m.get_default_chg_spin(), )[1] diff --git a/deepmd/dpmodel/atomic_model/make_base_atomic_model.py b/deepmd/dpmodel/atomic_model/make_base_atomic_model.py index 7118aa5d7a..f5deadde11 100644 --- a/deepmd/dpmodel/atomic_model/make_base_atomic_model.py +++ b/deepmd/dpmodel/atomic_model/make_base_atomic_model.py @@ -84,6 +84,14 @@ def get_nnei(self) -> int: """Returns the total number of selected neighboring atoms in the cut-off radius.""" return self.get_nsel() + def get_pair_exclude_types(self) -> list[tuple[int, int]]: + """Return the excluded atom-type pairs of this atomic model. + + Always set by ``__init__`` (empty list when no exclusion is + configured); an empty return means the pair-exclusion mask is off. + """ + return self.pair_exclude_types + @abstractmethod def get_dim_fparam(self) -> int: """Get the number (dimension) of frame parameters of this atomic model.""" diff --git a/deepmd/dpmodel/descriptor/descriptor.py b/deepmd/dpmodel/descriptor/descriptor.py index 1dce8fb7d1..d03fbe1a08 100644 --- a/deepmd/dpmodel/descriptor/descriptor.py +++ b/deepmd/dpmodel/descriptor/descriptor.py @@ -37,6 +37,12 @@ class DescriptorBlock(ABC, make_plugin_registry("DescriptorBlock")): local_cluster = False + # Stat-behavior flags with concrete defaults so stat machinery can read + # them on any block without getattr probes; blocks that configure them + # assign instance attributes in __init__ (issue #5897). + set_davg_zero: bool = False + set_stddev_constant: bool = False + def __new__(cls, *args: Any, **kwargs: Any) -> Any: if cls is DescriptorBlock: try: diff --git a/deepmd/dpmodel/descriptor/dpa1.py b/deepmd/dpmodel/descriptor/dpa1.py index 27d73b1251..1781728951 100644 --- a/deepmd/dpmodel/descriptor/dpa1.py +++ b/deepmd/dpmodel/descriptor/dpa1.py @@ -477,6 +477,10 @@ def __init__( self.concat_output_tebd = concat_output_tebd self.trainable = trainable self.precision = precision + # tebd-compression slots: declared here so presence is a class + # property, not a runtime accident (issue #5897); populated by + # enable_compression()/deserialize(). + self.type_embd_data = None self.tebd_compress = False self.geo_compress = False self.compress = False @@ -1113,6 +1117,10 @@ def enable_compression( stacklevel=2, ) + def get_geo_compress(self) -> bool: + """Return whether geometric tabulated compression is active.""" + return self.geo_compress + def serialize(self) -> dict: """Serialize the descriptor to dict.""" obj = self.se_atten @@ -1172,7 +1180,7 @@ def serialize(self) -> dict: if self.compress: type_embd_data = ( self.type_embd_data - if hasattr(self, "type_embd_data") + if self.type_embd_data is not None else obj.type_embd_data ) compress_dict: dict = { @@ -1514,6 +1522,9 @@ def __init__( self.mean = np.zeros(wanted_shape, dtype=PRECISION_DICT[self.precision]) self.stddev = np.ones(wanted_shape, dtype=PRECISION_DICT[self.precision]) self.orig_sel = self.sel + # tebd-compression slots: declared here so presence is a class + # property, not a runtime accident (issue #5897); populated by + # type_embedding_compression()/enable_compression(). self.tebd_compress = False self.geo_compress = False self.is_sorted = len(self.exclude_types) == 0 diff --git a/deepmd/dpmodel/descriptor/dpa2.py b/deepmd/dpmodel/descriptor/dpa2.py index d94bb5fc02..55566fee54 100644 --- a/deepmd/dpmodel/descriptor/dpa2.py +++ b/deepmd/dpmodel/descriptor/dpa2.py @@ -1421,6 +1421,10 @@ def _call_dense( g1 = xp.concat([g1, g1_inp], axis=-1) return g1, rot_mat, g2, h2, sw + def get_geo_compress(self) -> bool: + """Return whether geometric tabulated compression is active.""" + return self.geo_compress + def serialize(self) -> dict: repinit = self.repinit repformers = self.repformers diff --git a/deepmd/dpmodel/descriptor/dpa3.py b/deepmd/dpmodel/descriptor/dpa3.py index aa1d10fc33..111dc7a520 100644 --- a/deepmd/dpmodel/descriptor/dpa3.py +++ b/deepmd/dpmodel/descriptor/dpa3.py @@ -543,10 +543,6 @@ def get_dim_chg_spin(self) -> int: """Returns the dimension of charge_spin input.""" return 2 if self.add_chg_spin_ebd else 0 - def has_default_chg_spin(self) -> bool: - """Returns whether default charge_spin values are set.""" - return self.default_chg_spin is not None - def get_default_chg_spin(self) -> list[float] | None: """Returns the default charge_spin values.""" return self.default_chg_spin diff --git a/deepmd/dpmodel/descriptor/dpa4.py b/deepmd/dpmodel/descriptor/dpa4.py index 801b4c34f2..cce49335dc 100644 --- a/deepmd/dpmodel/descriptor/dpa4.py +++ b/deepmd/dpmodel/descriptor/dpa4.py @@ -2286,10 +2286,6 @@ def get_dim_chg_spin(self) -> int: """Return the charge/spin condition width.""" return 2 if self.add_chg_spin_ebd else 0 - def has_default_chg_spin(self) -> bool: - """Return whether default charge/spin conditions are configured.""" - return self.default_chg_spin is not None - def get_default_chg_spin(self) -> list[float] | None: """Return default charge/spin conditions.""" return self.default_chg_spin diff --git a/deepmd/dpmodel/descriptor/hybrid.py b/deepmd/dpmodel/descriptor/hybrid.py index 9709ce37f0..83a903c687 100644 --- a/deepmd/dpmodel/descriptor/hybrid.py +++ b/deepmd/dpmodel/descriptor/hybrid.py @@ -129,33 +129,26 @@ def get_dim_chg_spin(self) -> int: (descrpt.get_dim_chg_spin() for descrpt in self.descrpt_list), default=0 ) - def has_default_chg_spin(self) -> bool: - """Returns whether the descriptor has a default charge_spin value.""" + def get_default_chg_spin(self) -> list[float] | None: + """Returns the default charge_spin value, or None. + + ``None`` unless every sub-descriptor that supports charge_spin + (``get_dim_chg_spin() > 0``) agrees on the same default value. + """ default_chg_spin = None found_chg_spin = False for descrpt in self.descrpt_list: if descrpt.get_dim_chg_spin() == 0: continue found_chg_spin = True - if not descrpt.has_default_chg_spin(): - return False child_default_chg_spin = descrpt.get_default_chg_spin() if child_default_chg_spin is None: - return False + return None if default_chg_spin is None: default_chg_spin = child_default_chg_spin elif child_default_chg_spin != default_chg_spin: - return False - return found_chg_spin - - def get_default_chg_spin(self) -> list[float] | None: - """Returns the default charge_spin value, or None.""" - if not self.has_default_chg_spin(): - return None - for descrpt in self.descrpt_list: - if descrpt.get_dim_chg_spin() > 0: - return descrpt.get_default_chg_spin() - return None + return None + return default_chg_spin if found_chg_spin else None def get_rcut_smth(self) -> float: """Returns the radius where the neighbor information starts to smoothly decay to 0.""" diff --git a/deepmd/dpmodel/descriptor/make_base_descriptor.py b/deepmd/dpmodel/descriptor/make_base_descriptor.py index 2d058a3398..b0a77e22fc 100644 --- a/deepmd/dpmodel/descriptor/make_base_descriptor.py +++ b/deepmd/dpmodel/descriptor/make_base_descriptor.py @@ -48,6 +48,15 @@ def make_base_descriptor( class BD(ABC, PluginVariant, make_plugin_registry("descriptor")): """Base descriptor provides the interfaces of descriptor.""" + # Stat-behavior flags with concrete defaults so stat machinery (e.g. + # ``merge_env_stat``, which accepts either a ``Descriptor`` or a + # ``DescriptorBlock``) can read them on any descriptor without + # getattr probes; descriptors that configure them assign instance + # attributes in __init__ (issue #5897). Mirrors the same defaults on + # ``DescriptorBlock``. + set_davg_zero: bool = False + set_stddev_constant: bool = False + def __new__(cls, *args: Any, **kwargs: Any) -> Any: if cls is BD: cls = cls.get_class_by_type(j_get_type(kwargs, cls.__name__)) @@ -100,14 +109,18 @@ def get_dim_chg_spin(self) -> int: """Returns the dimension of charge_spin input (0 if not supported).""" return 0 - def has_default_chg_spin(self) -> bool: - """Returns whether the descriptor has a default charge_spin value.""" - return False - def get_default_chg_spin(self) -> Any: """Returns the default charge_spin value, or None.""" return None + def get_geo_compress(self) -> bool: + """Return whether geometric tabulated compression is active. + + Concrete default ``False``; descriptor families with a + geometric compression path override from their own state. + """ + return False + @abstractmethod def mixed_types(self) -> bool: """Returns if the descriptor requires a neighbor list that distinguish different diff --git a/deepmd/dpmodel/descriptor/se_atten_v2.py b/deepmd/dpmodel/descriptor/se_atten_v2.py index ec7e3e4f77..9a7d8a82fe 100644 --- a/deepmd/dpmodel/descriptor/se_atten_v2.py +++ b/deepmd/dpmodel/descriptor/se_atten_v2.py @@ -259,7 +259,7 @@ def serialize(self) -> dict: if self.compress: type_embd_data = ( self.type_embd_data - if hasattr(self, "type_embd_data") + if self.type_embd_data is not None else obj.type_embd_data ) compress_dict: dict = { diff --git a/deepmd/dpmodel/descriptor/se_t_tebd.py b/deepmd/dpmodel/descriptor/se_t_tebd.py index 3d5321c0ef..b810334031 100644 --- a/deepmd/dpmodel/descriptor/se_t_tebd.py +++ b/deepmd/dpmodel/descriptor/se_t_tebd.py @@ -206,6 +206,11 @@ def __init__( self.trainable = trainable self.precision = precision self.compress = False + # tebd-compression slot: declared here so presence is a class + # property, not a runtime accident (issue #5897); optionally + # populated by deserialize() (only present in compressed models + # that carry type embedding compression data). + self.type_embd_data = None def get_rcut(self) -> float: """Returns the cut-off radius.""" @@ -463,7 +468,7 @@ def serialize(self) -> dict: "compress_info": [to_numpy_array(i) for i in self.compress_info], }, } - if hasattr(self, "type_embd_data"): + if self.type_embd_data is not None: compress_dict["@variables"]["type_embd_data"] = to_numpy_array( self.type_embd_data ) diff --git a/deepmd/dpmodel/fitting/make_base_fitting.py b/deepmd/dpmodel/fitting/make_base_fitting.py index cf8172bd03..7a595d4bc8 100644 --- a/deepmd/dpmodel/fitting/make_base_fitting.py +++ b/deepmd/dpmodel/fitting/make_base_fitting.py @@ -67,6 +67,30 @@ def compute_output_stats(self, merged: Any) -> NoReturn: """Update the output bias for fitting net.""" raise NotImplementedError + def reinit_exclude(self, exclude_types: list[int] = []) -> None: + """Reinitialize the per-type output exclusion list. + + Concrete default for fittings without exclusion support: an + empty list is a no-op; a non-empty list raises, because + silently ignoring a requested exclusion would degrade the + model without any signal. + + Parameters + ---------- + exclude_types + Atom types whose fitting output is excluded. + + Raises + ------ + NotImplementedError + If ``exclude_types`` is non-empty and this fitting does + not support atom-type exclusion. + """ + if exclude_types: + raise NotImplementedError( + "this fitting does not support atom-type exclusion" + ) + @abstractmethod def get_type_map(self) -> list[str]: """Get the name to each type of atoms.""" diff --git a/deepmd/dpmodel/model/base_model.py b/deepmd/dpmodel/model/base_model.py index 9c85fa7e26..0ee9763b1f 100644 --- a/deepmd/dpmodel/model/base_model.py +++ b/deepmd/dpmodel/model/base_model.py @@ -107,6 +107,45 @@ def has_spin(self) -> bool: """ return False + def has_chg_spin_ebd(self) -> bool: + """Return whether the model conditions on charge/spin embedding. + + Concrete default ``False``; models wrapping an atomic model + override to delegate. + """ + return False + + def get_dim_chg_spin(self) -> int: + """Return the charge/spin condition width (0 if unsupported).""" + return 0 + + def get_default_chg_spin(self) -> list | None: + """Return default charge/spin conditions, or ``None`` if none + are configured. ``is not None`` is the support predicate. + """ + return None + + def get_var_name(self) -> str | None: + """Return the fitted property's variable name, or ``None`` if + this is not a property model. ``is not None`` is the support + predicate. + """ + return None + + def get_task_dim(self) -> int: + """Return the property output dimension (property models only). + + Raises + ------ + NotImplementedError + If the model is not a property model. + """ + raise NotImplementedError("get_task_dim: property models only") + + def get_intensive(self) -> bool: + """Return whether the fitted property is intensive.""" + return False + @abstractmethod def serialize(self) -> dict: """Serialize the model. diff --git a/deepmd/dpmodel/model/make_model.py b/deepmd/dpmodel/model/make_model.py index 44967bba74..73b53ce6ea 100644 --- a/deepmd/dpmodel/model/make_model.py +++ b/deepmd/dpmodel/model/make_model.py @@ -428,7 +428,7 @@ def call_common( charge_spin=cs, neighbor_list=neighbor_list, # exclusion is a nlist-BUILD transform (decision #18/A4) - pair_excl=getattr(self.atomic_model, "pair_excl", None), + pair_excl=self.atomic_model.pair_excl, ) model_predict = self._output_type_cast(model_predict, input_prec) return model_predict @@ -527,7 +527,7 @@ def _call_common_graph( # is constructed, so the graph lower / exported ``.pt2`` consumes an # already-excluded ``edge_mask`` and never re-applies it. Mirrors the # pt_expt eager path and the C++ ``applyPairExclusion`` at build. - pair_excl = getattr(self.atomic_model, "pair_excl", None) + pair_excl = self.atomic_model.pair_excl if method == "dense": ng = build_neighbor_graph( cc, atype, bb, self.get_rcut(), pair_excl=pair_excl @@ -1131,10 +1131,6 @@ def get_dim_chg_spin(self) -> int: """Get the dimension of charge_spin input.""" return self.atomic_model.get_dim_chg_spin() - def has_default_chg_spin(self) -> bool: - """Check if the model has default charge_spin values.""" - return self.atomic_model.has_default_chg_spin() - def get_default_chg_spin(self) -> list[float] | None: """Get the default charge_spin values.""" return self.atomic_model.get_default_chg_spin() diff --git a/deepmd/dpmodel/model/spin_model.py b/deepmd/dpmodel/model/spin_model.py index 8744cc7b9f..ed3d9e3cc2 100644 --- a/deepmd/dpmodel/model/spin_model.py +++ b/deepmd/dpmodel/model/spin_model.py @@ -210,12 +210,16 @@ def process_spin_input_lower( # pair exclusion in here (decision #18/A4 — the lower consumes a # pre-excluded nlist and never re-applies it). No-op when the backbone # has no pair_exclude_types. - pair_excl = getattr( - self.backbone_model.atomic_model + # ``backbone_model`` is either a full ``make_model``-wrapped model + # (exposes ``.atomic_model``) or, per the ``__init__`` annotation, + # a bare ``DPAtomicModel`` -- which always carries ``pair_excl`` + # (set unconditionally by ``BaseAtomicModel.__init__`` via + # ``reinit_pair_exclude``, ``None`` when no exclusion is + # configured). Direct access, not a defensive ``getattr`` probe. + pair_excl = ( + self.backbone_model.atomic_model.pair_excl if hasattr(self.backbone_model, "atomic_model") - else self.backbone_model, - "pair_excl", - None, + else self.backbone_model.pair_excl ) if pair_excl is not None: from deepmd.dpmodel.utils.nlist import ( diff --git a/deepmd/dpmodel/utils/env_mat_stat.py b/deepmd/dpmodel/utils/env_mat_stat.py index 8288ede2e7..1bb123534d 100644 --- a/deepmd/dpmodel/utils/env_mat_stat.py +++ b/deepmd/dpmodel/utils/env_mat_stat.py @@ -69,9 +69,7 @@ def merge_env_stat( or getattr(link_obj, "stats", None) is None ): return - if getattr(base_obj, "set_stddev_constant", False) and getattr( - base_obj, "set_davg_zero", False - ): + if base_obj.set_stddev_constant and base_obj.set_davg_zero: return # Weighted merge of StatItem objects @@ -93,7 +91,7 @@ def merge_env_stat( xp = array_api_compat.array_namespace(current_stddev) device = array_api_compat.device(current_stddev) merged_mean = current_mean - if not getattr(base_obj, "set_davg_zero", False): + if not base_obj.set_davg_zero: merged_mean = xp.asarray(mean, dtype=current_mean.dtype, device=device) merged_stddev = xp.asarray( stddev, diff --git a/deepmd/infer/deep_eval.py b/deepmd/infer/deep_eval.py index 915d216672..d20626e835 100644 --- a/deepmd/infer/deep_eval.py +++ b/deepmd/infer/deep_eval.py @@ -390,10 +390,11 @@ def _get_property_var_name(model: Any) -> str | None: """Return the property variable name of ``model``, or ``None``. Used by every backend's ``model_type`` to detect a property model. - ``get_var_name`` may be absent (dpmodel/pt live models expose it only on - property models) or present-but-unimplemented (jax/tf2 artifacts always - define it and raise ``NotImplementedError`` otherwise), so probe - defensively. + Live dpmodel-family models always expose ``get_var_name`` (the base + class's concrete default returns ``None`` for non-property models); + ``get_var_name`` is only absent on frozen pt legacy models, and only + present-but-unimplemented -- raising ``NotImplementedError`` -- on + jax/tf2 artifacts, so probe defensively. """ if not hasattr(model, "get_var_name"): return None @@ -441,6 +442,7 @@ def get_has_hessian(self) -> bool: def get_var_name(self) -> str: """Get the name of the fitting property (property models only).""" model = self.get_model() + # artifact boundary: get_model() may return a loaded SavedModel if hasattr(model, "get_var_name"): return model.get_var_name() raise NotImplementedError @@ -448,6 +450,7 @@ def get_var_name(self) -> str: def get_task_dim(self) -> int: """Get the output dimension of the property (property models only).""" model = self.get_model() + # artifact boundary: get_model() may return a loaded SavedModel if hasattr(model, "get_task_dim"): return model.get_task_dim() raise NotImplementedError @@ -455,6 +458,7 @@ def get_task_dim(self) -> int: def get_intensive(self) -> bool: """Whether the property is intensive (property models only).""" model = self.get_model() + # artifact boundary: get_model() may return a loaded SavedModel if hasattr(model, "get_intensive"): return model.get_intensive() raise NotImplementedError diff --git a/deepmd/jax/infer/deep_eval.py b/deepmd/jax/infer/deep_eval.py index 0e6c11ede6..d2930ff8c7 100644 --- a/deepmd/jax/infer/deep_eval.py +++ b/deepmd/jax/infer/deep_eval.py @@ -524,18 +524,17 @@ def has_default_fparam(self) -> bool: def has_chg_spin_ebd(self) -> bool: """Check if the model has charge spin embedding.""" - if hasattr(self.dp, "has_chg_spin_ebd"): - return self.dp.has_chg_spin_ebd() - return False + return self.dp.has_chg_spin_ebd() def get_dim_chg_spin(self) -> int: """Get the dimension of charge_spin input.""" - if hasattr(self.dp, "get_dim_chg_spin"): - return self.dp.get_dim_chg_spin() - return 0 + return self.dp.get_dim_chg_spin() def has_default_chg_spin(self) -> bool: - """Check if the model has default charge_spin values.""" - if hasattr(self.dp, "has_default_chg_spin"): - return self.dp.has_default_chg_spin() - return False + """Check if the model has default charge_spin values. + + ``has_default_chg_spin`` was merged into ``get_default_chg_spin`` on + the live-model interfaces; this wrapper method is kept for API + stability and computes the predicate directly. + """ + return self.dp.get_default_chg_spin() is not None diff --git a/deepmd/jax/jax2tf/serialization.py b/deepmd/jax/jax2tf/serialization.py index 24ea448843..170cc516f5 100644 --- a/deepmd/jax/jax2tf/serialization.py +++ b/deepmd/jax/jax2tf/serialization.py @@ -264,6 +264,10 @@ def call( aparam: tf.Tensor | None = None, charge_spin: tf.Tensor | None = None, ) -> dict[str, tf.Tensor]: + # exclusion is a nlist-BUILD transform (decision #18/A4); the + # traced lower consumes a pre-excluded nlist. Guard + # atomic_model too: test doubles (DummyModel) lack it. + am = getattr(model, "atomic_model", None) return model_call_from_call_lower( call_lower=call_lower, rcut=model.get_rcut(), @@ -277,12 +281,7 @@ def call( aparam=aparam, charge_spin=charge_spin, do_atomic_virial=do_atomic_virial, - # exclusion is a nlist-BUILD transform (decision #18/A4); - # the traced lower consumes a pre-excluded nlist. Guard - # atomic_model too: test doubles (DummyModel) lack it. - pair_excl=getattr( - getattr(model, "atomic_model", None), "pair_excl", None - ), + pair_excl=am.pair_excl if am is not None else None, ) return call @@ -418,9 +417,8 @@ def get_pair_exclude_types() -> tf.Tensor: # the LAMMPS nlist before the traced call_lower_* consumes it # (decision #18/A4). The upper ``call`` already pre-excludes its # freshly built nlist. Guard atomic_model: test doubles may lack it. - pet = getattr( - getattr(model, "atomic_model", None), "pair_exclude_types", [] - ) + am = getattr(model, "atomic_model", None) + pet = am.get_pair_exclude_types() if am is not None else [] flat = [int(t) for pair in (pet or []) for t in pair] return tf.constant(flat, dtype=tf.int64) @@ -470,7 +468,7 @@ def get_dim_chg_spin() -> tf.Tensor: @tf.function def has_default_chg_spin() -> tf.Tensor: - return tf.constant(model.has_default_chg_spin(), dtype=tf.bool) + return tf.constant(model.get_default_chg_spin() is not None, dtype=tf.bool) tf_model.has_default_chg_spin = has_default_chg_spin @@ -485,7 +483,8 @@ def get_default_chg_spin() -> tf.Tensor: # property models: persist the output name/dimension/intensiveness so # the evaluator can dispatch to DeepProperty and reshape the output. - if hasattr(model, "get_var_name"): + is_property_model = model.get_var_name() is not None + if is_property_model: @tf.function def get_var_name() -> tf.Tensor: diff --git a/deepmd/jax/jax2tf/tfmodel.py b/deepmd/jax/jax2tf/tfmodel.py index 57efdb5c14..6347af2fc6 100644 --- a/deepmd/jax/jax2tf/tfmodel.py +++ b/deepmd/jax/jax2tf/tfmodel.py @@ -41,6 +41,32 @@ def decode_list_of_bytes(list_of_bytes: list[bytes]) -> list[str]: return [x.decode() for x in list_of_bytes] +def _decode_default_chg_spin( + has_default: bool, default_chg_spin: list[float] +) -> list[float] | None: + """Decode the exported default charge-spin values. + + The exporter always writes a ``get_default_chg_spin`` tensor next to + the boolean marker, using an EMPTY tensor when the model has no + default. Normalize the no-default case to ``None`` so the wrapper + upholds the live-model invariant + ``get_default_chg_spin() is None == no default``. + + Parameters + ---------- + has_default : bool + The boolean marker exported by ``has_default_chg_spin``. + default_chg_spin : list[float] + The decoded ``get_default_chg_spin`` tensor values. + + Returns + ------- + list[float] | None + The default values, or ``None`` when the artifact has none. + """ + return default_chg_spin if has_default else None + + class TFModelWrapper(tf.Module): def __init__( self, @@ -93,10 +119,13 @@ def __init__( if hasattr(self.model, "has_default_chg_spin") else False ) - self.default_chg_spin = ( - self.model.get_default_chg_spin().numpy().tolist() - if hasattr(self.model, "get_default_chg_spin") - else None + self.default_chg_spin = _decode_default_chg_spin( + self._has_default_chg_spin, + ( + self.model.get_default_chg_spin().numpy().tolist() + if hasattr(self.model, "get_default_chg_spin") + else [] + ), ) # property models only (absent for other model types). if hasattr(self.model, "get_var_name"): diff --git a/deepmd/jax/jax_md/__init__.py b/deepmd/jax/jax_md/__init__.py index 8450b96e5b..fa3859d6dc 100644 --- a/deepmd/jax/jax_md/__init__.py +++ b/deepmd/jax/jax_md/__init__.py @@ -328,7 +328,8 @@ def _eval_with_jax_md_neighbor( # re-applies it. The JAX-MD neighbor list is built without exclusion, so # fold it in at this ingestion seam -- otherwise excluded pairs would be # silently included (fail-open). - pair_excl = getattr(getattr(model, "atomic_model", None), "pair_excl", None) + am = getattr(model, "atomic_model", None) + pair_excl = am.pair_excl if am is not None else None if pair_excl is not None: from deepmd.dpmodel.utils.nlist import ( apply_pair_exclusion_nlist, diff --git a/deepmd/jax/train/trainer.py b/deepmd/jax/train/trainer.py index d78dbc38ed..a35de2a49b 100644 --- a/deepmd/jax/train/trainer.py +++ b/deepmd/jax/train/trainer.py @@ -973,7 +973,7 @@ def _prepare_batch( box=jax_data["box"] if jax_data["find_box"] else None, fparam=jax_data.get("fparam", None), aparam=jax_data.get("aparam", None), - pair_excl=getattr(model.atomic_model, "pair_excl", None), + pair_excl=model.atomic_model.pair_excl, ) return jax_data, extended_coord, extended_atype, nlist, mapping, fp, ap diff --git a/deepmd/jax/utils/serialization.py b/deepmd/jax/utils/serialization.py index a347c2dff8..62a6851160 100644 --- a/deepmd/jax/utils/serialization.py +++ b/deepmd/jax/utils/serialization.py @@ -335,6 +335,7 @@ def call_lower_with_fixed_do_atomic_virial( data["@variables"]["stablehlo_atomic_virial_no_ghost"] = np.void( serialized_atomic_virial_no_ghost ) + is_property_model = model.get_var_name() is not None data["constants"] = { "type_map": model.get_type_map(), "rcut": model.get_rcut(), @@ -352,15 +353,9 @@ def call_lower_with_fixed_do_atomic_virial( # property models: the output name/dimension/intensiveness cannot be # recovered from the StableHLO alone, so persist them for the # evaluator (None for non-property models). - "var_name": model.get_var_name() - if hasattr(model, "get_var_name") - else None, - "task_dim": model.get_task_dim() - if hasattr(model, "get_task_dim") - else None, - "intensive": model.get_intensive() - if hasattr(model, "get_intensive") - else False, + "var_name": model.get_var_name(), + "task_dim": model.get_task_dim() if is_property_model else None, + "intensive": model.get_intensive() if is_property_model else False, } save_dp_model(filename=model_file, model_dict=data) elif model_file.endswith(".savedmodel"): diff --git a/deepmd/kernels/cuda/dpa1/canonical.py b/deepmd/kernels/cuda/dpa1/canonical.py index ebfdd9981d..856fab953c 100644 --- a/deepmd/kernels/cuda/dpa1/canonical.py +++ b/deepmd/kernels/cuda/dpa1/canonical.py @@ -34,12 +34,12 @@ def canonical_model_eligible(model: Any) -> bool: fitting = getattr(atomic_model, "fitting_net", None) if descriptor is None or fitting is None: return False - if not bool(getattr(descriptor, "geo_compress", False)): + if not descriptor.get_geo_compress(): return False eligible = getattr(descriptor, "_fused_eligible", None) if not callable(eligible) or not bool(eligible("cuda")): return False - if getattr(atomic_model, "pair_excl", None) is not None: + if atomic_model.pair_excl is not None: return False if getattr(atomic_model, "atom_excl", None) is not None: return False diff --git a/deepmd/pd/infer/deep_eval.py b/deepmd/pd/infer/deep_eval.py index 749bcdb76b..fd393d901d 100644 --- a/deepmd/pd/infer/deep_eval.py +++ b/deepmd/pd/infer/deep_eval.py @@ -239,12 +239,10 @@ def get_intensive(self) -> bool: def get_var_name(self) -> str: """Get the name of the property.""" - if hasattr(self.dp.model["Default"], "get_var_name") and callable( - getattr(self.dp.model["Default"], "get_var_name") - ): - return self.dp.model["Default"].get_var_name() - else: + var_name = self.dp.model["Default"].get_var_name() + if var_name is None: raise NotImplementedError + return var_name @property def model_type(self) -> type["DeepEvalWrapper"]: diff --git a/deepmd/pd/model/atomic_model/base_atomic_model.py b/deepmd/pd/model/atomic_model/base_atomic_model.py index 50fc4b4a5f..5a1ff9d2cd 100644 --- a/deepmd/pd/model/atomic_model/base_atomic_model.py +++ b/deepmd/pd/model/atomic_model/base_atomic_model.py @@ -189,10 +189,6 @@ def get_dim_chg_spin(self) -> int: """Get the dimension of charge_spin input.""" return 0 - def has_default_chg_spin(self) -> bool: - """Check if the model has default charge_spin values.""" - return False - def get_default_chg_spin(self) -> paddle.Tensor | None: """Get the default charge_spin values.""" return None diff --git a/deepmd/pd/model/atomic_model/dp_atomic_model.py b/deepmd/pd/model/atomic_model/dp_atomic_model.py index 1b37184b18..334ebdc5f5 100644 --- a/deepmd/pd/model/atomic_model/dp_atomic_model.py +++ b/deepmd/pd/model/atomic_model/dp_atomic_model.py @@ -65,9 +65,7 @@ def __init__( self.sel = self.descriptor.get_sel() self.fitting_net = fitting super().init_out_stat() - self.add_chg_spin_ebd: bool = getattr( - self.descriptor, "add_chg_spin_ebd", False - ) + self.add_chg_spin_ebd: bool = self.descriptor.get_dim_chg_spin() > 0 self.enable_eval_descriptor_hook = False self.enable_eval_fitting_last_layer_hook = False self.eval_descriptor_list = [] @@ -85,13 +83,12 @@ def _string_to_array(s: str | list[str]) -> list[int]: ), ) self.buffer_type_map.name = "buffer_type_map" - if hasattr(self.descriptor, "has_message_passing"): - # register 'has_message_passing' as buffer(cast to int32 as problems may meets with vector) - self.register_buffer( - "buffer_has_message_passing", - paddle.to_tensor(self.descriptor.has_message_passing(), dtype="int32"), - ) - self.buffer_has_message_passing.name = "buffer_has_message_passing" + # register 'has_message_passing' as buffer(cast to int32 as problems may meets with vector) + self.register_buffer( + "buffer_has_message_passing", + paddle.to_tensor(self.descriptor.has_message_passing(), dtype="int32"), + ) + self.buffer_has_message_passing.name = "buffer_has_message_passing" # register 'ntypes' as buffer self.register_buffer( "buffer_ntypes", paddle.to_tensor(self.ntypes, dtype="int32") @@ -488,14 +485,8 @@ def get_dim_chg_spin(self) -> int: return self.descriptor.get_dim_chg_spin() return 0 - def has_default_chg_spin(self) -> bool: - """Check if the model has default charge_spin values.""" - if self.add_chg_spin_ebd: - return self.descriptor.has_default_chg_spin() - return False - def get_default_chg_spin(self) -> paddle.Tensor | None: """Get the default charge_spin values as a tensor.""" - if self.add_chg_spin_ebd and self.descriptor.has_default_chg_spin(): + if self.add_chg_spin_ebd: return self.descriptor.get_default_chg_spin() return None diff --git a/deepmd/pd/model/descriptor/descriptor.py b/deepmd/pd/model/descriptor/descriptor.py index 6071ab7d2b..824adaad8d 100644 --- a/deepmd/pd/model/descriptor/descriptor.py +++ b/deepmd/pd/model/descriptor/descriptor.py @@ -47,6 +47,12 @@ class DescriptorBlock(paddle.nn.Layer, ABC, make_plugin_registry("DescriptorBloc local_cluster = False + # Stat-behavior flags with concrete defaults so stat machinery can read + # them on any block without getattr probes; blocks that configure them + # assign instance attributes in __init__ (issue #5897). + set_davg_zero: bool = False + set_stddev_constant: bool = False + def __new__(cls, *args: Any, **kwargs: Any) -> Self: if cls is DescriptorBlock: try: @@ -145,8 +151,7 @@ def share_params( # link buffers if hasattr(self, "mean"): if not resume and ( - not getattr(self, "set_stddev_constant", False) - or not getattr(self, "set_davg_zero", False) + not self.set_stddev_constant or not self.set_davg_zero ): # in case of change params during resume base_env = EnvMatStatSe(base_class) diff --git a/deepmd/pd/model/descriptor/dpa1.py b/deepmd/pd/model/descriptor/dpa1.py index ccd8944cc8..8db2792117 100644 --- a/deepmd/pd/model/descriptor/dpa1.py +++ b/deepmd/pd/model/descriptor/dpa1.py @@ -370,10 +370,6 @@ def get_dim_chg_spin(self) -> int: """Returns the dimension of charge_spin input (0 if not supported).""" return 0 - def has_default_chg_spin(self) -> bool: - """Returns whether the descriptor has a default charge_spin value.""" - return False - def get_default_chg_spin(self) -> None: """Returns the default charge_spin value, or None.""" return None diff --git a/deepmd/pd/model/descriptor/dpa2.py b/deepmd/pd/model/descriptor/dpa2.py index b976ed6dd5..8ca36e2091 100644 --- a/deepmd/pd/model/descriptor/dpa2.py +++ b/deepmd/pd/model/descriptor/dpa2.py @@ -337,10 +337,6 @@ def get_dim_chg_spin(self) -> int: """Returns the dimension of charge_spin input (0 if not supported).""" return 0 - def has_default_chg_spin(self) -> bool: - """Returns whether the descriptor has a default charge_spin value.""" - return False - def get_default_chg_spin(self) -> None: """Returns the default charge_spin value, or None.""" return None diff --git a/deepmd/pd/model/descriptor/dpa3.py b/deepmd/pd/model/descriptor/dpa3.py index 95fe69422b..9a83fe56c4 100644 --- a/deepmd/pd/model/descriptor/dpa3.py +++ b/deepmd/pd/model/descriptor/dpa3.py @@ -458,10 +458,6 @@ def get_dim_chg_spin(self) -> int: """Returns the dimension of charge_spin input.""" return 2 if self.add_chg_spin_ebd else 0 - def has_default_chg_spin(self) -> bool: - """Returns whether default charge_spin values are set.""" - return self.default_chg_spin is not None - def get_default_chg_spin(self) -> paddle.Tensor | None: """Get the default charge_spin values as a tensor.""" if self.default_chg_spin is None: diff --git a/deepmd/pd/model/descriptor/se_a.py b/deepmd/pd/model/descriptor/se_a.py index 0098ca2186..918b8be224 100644 --- a/deepmd/pd/model/descriptor/se_a.py +++ b/deepmd/pd/model/descriptor/se_a.py @@ -123,10 +123,6 @@ def get_dim_chg_spin(self) -> int: """Returns the dimension of charge_spin input (0 if not supported).""" return 0 - def has_default_chg_spin(self) -> bool: - """Returns whether the descriptor has a default charge_spin value.""" - return False - def get_default_chg_spin(self) -> None: """Returns the default charge_spin value, or None.""" return None diff --git a/deepmd/pd/model/descriptor/se_t_tebd.py b/deepmd/pd/model/descriptor/se_t_tebd.py index d61294650a..c83e22561d 100644 --- a/deepmd/pd/model/descriptor/se_t_tebd.py +++ b/deepmd/pd/model/descriptor/se_t_tebd.py @@ -194,10 +194,6 @@ def get_dim_chg_spin(self) -> int: """Returns the dimension of charge_spin input (0 if not supported).""" return 0 - def has_default_chg_spin(self) -> bool: - """Returns whether the descriptor has a default charge_spin value.""" - return False - def get_default_chg_spin(self) -> None: """Returns the default charge_spin value, or None.""" return None diff --git a/deepmd/pd/model/model/make_model.py b/deepmd/pd/model/model/make_model.py index 7db800bf22..f5beb501fc 100644 --- a/deepmd/pd/model/model/make_model.py +++ b/deepmd/pd/model/model/make_model.py @@ -562,10 +562,6 @@ def get_dim_chg_spin(self) -> int: """Get the dimension of charge_spin input.""" return self.atomic_model.get_dim_chg_spin() - def has_default_chg_spin(self) -> bool: - """Check if the model has default charge_spin values.""" - return self.atomic_model.has_default_chg_spin() - def get_default_chg_spin(self) -> paddle.Tensor | None: """Get the default charge_spin values.""" return self.atomic_model.get_default_chg_spin() diff --git a/deepmd/pd/train/training.py b/deepmd/pd/train/training.py index 53f768fded..688f0c81c0 100644 --- a/deepmd/pd/train/training.py +++ b/deepmd/pd/train/training.py @@ -1336,10 +1336,9 @@ def get_additional_data_requirement(_model: Any) -> list[DataRequirementItem]: ] additional_data_requirement += spin_requirement_items if _model.has_chg_spin_ebd(): - has_default_cs = _model.has_default_chg_spin() - cs_default = ( - _model.get_default_chg_spin().cpu().numpy() if has_default_cs else 0.0 - ) + default_cs = _model.get_default_chg_spin() + has_default_cs = default_cs is not None + cs_default = default_cs.cpu().numpy() if has_default_cs else 0.0 additional_data_requirement.append( DataRequirementItem( "charge_spin", diff --git a/deepmd/pt_expt/common.py b/deepmd/pt_expt/common.py index 5676652ba3..dbff99e3bd 100644 --- a/deepmd/pt_expt/common.py +++ b/deepmd/pt_expt/common.py @@ -329,6 +329,35 @@ def dpmodel_setattr(obj: torch.nn.Module, name: str, value: Any) -> tuple[bool, return False, value +def register_buffer_replacing_slot( + obj: torch.nn.Module, name: str, tensor: torch.Tensor +) -> None: + """Register a buffer, replacing a same-named plain attribute if present. + + Some dpmodel base ``__init__``s declare a capability slot such as + ``self.type_embd_data = None`` so its presence is a class property + rather than a runtime accident (issue #5897). A torch-native + compression path may later want to store the real value as a + persistent buffer via ``torch.nn.Module.register_buffer`` directly + (bypassing ``__setattr__``/``dpmodel_setattr``, e.g. because the + value is already a computed ``torch.Tensor`` rather than a + ``np.ndarray``). ``register_buffer`` raises ``KeyError`` if the name + already exists as a plain (non-buffer) attribute, so remove it first. + + Parameters + ---------- + obj : torch.nn.Module + The pt_expt wrapper object to register the buffer on. + name : str + The buffer name. + tensor : torch.Tensor + The tensor value to store as a buffer. + """ + if hasattr(obj, name) and name not in obj._buffers: + delattr(obj, name) + torch.nn.Module.register_buffer(obj, name, tensor) + + # --------------------------------------------------------------------------- # Utility # --------------------------------------------------------------------------- diff --git a/deepmd/pt_expt/descriptor/dpa1.py b/deepmd/pt_expt/descriptor/dpa1.py index 47145d96ec..91266b8788 100644 --- a/deepmd/pt_expt/descriptor/dpa1.py +++ b/deepmd/pt_expt/descriptor/dpa1.py @@ -56,6 +56,7 @@ triton_infer_level, ) from deepmd.pt_expt.common import ( + register_buffer_replacing_slot, torch_module, ) from deepmd.pt_expt.descriptor.base_descriptor import ( @@ -529,7 +530,7 @@ def _store_type_embd_data(self) -> None: self.se_atten.embeddings_strip[0].call(two_side_embd).detach() ) - torch.nn.Module.register_buffer(self, "type_embd_data", embd_tensor) + register_buffer_replacing_slot(self, "type_embd_data", embd_tensor) @cast_precision def call( diff --git a/deepmd/pt_expt/descriptor/dpa2.py b/deepmd/pt_expt/descriptor/dpa2.py index 6d94f5f5cc..46bc774cc9 100644 --- a/deepmd/pt_expt/descriptor/dpa2.py +++ b/deepmd/pt_expt/descriptor/dpa2.py @@ -21,6 +21,7 @@ remap_atype_to_padding, ) from deepmd.pt_expt.common import ( + register_buffer_replacing_slot, torch_module, ) from deepmd.pt_expt.descriptor.base_descriptor import ( @@ -284,7 +285,7 @@ def _store_type_embd_data(self) -> None: self.repinit.embeddings_strip[0].call(two_side_embd).detach() ) - torch.nn.Module.register_buffer(self, "type_embd_data", embd_tensor) + register_buffer_replacing_slot(self, "type_embd_data", embd_tensor) @cast_precision def call( diff --git a/deepmd/pt_expt/descriptor/repflows.py b/deepmd/pt_expt/descriptor/repflows.py index dacab9f464..2b88da3562 100644 --- a/deepmd/pt_expt/descriptor/repflows.py +++ b/deepmd/pt_expt/descriptor/repflows.py @@ -57,7 +57,7 @@ def _exchange_ghosts( # entirely, so combining it with comm_dict is contradictory. # Surface this as a clear error rather than producing silently # wrong results. - if getattr(self, "use_loc_mapping", False): + if self.use_loc_mapping: raise RuntimeError( "DescrptBlockRepflows._exchange_ghosts: comm_dict is " "set but use_loc_mapping=True. Multi-rank parallel " diff --git a/deepmd/pt_expt/descriptor/se_t_tebd.py b/deepmd/pt_expt/descriptor/se_t_tebd.py index 512e0ebe7a..69b786552d 100644 --- a/deepmd/pt_expt/descriptor/se_t_tebd.py +++ b/deepmd/pt_expt/descriptor/se_t_tebd.py @@ -16,6 +16,7 @@ remap_atype_to_padding, ) from deepmd.pt_expt.common import ( + register_buffer_replacing_slot, torch_module, ) from deepmd.pt_expt.descriptor.base_descriptor import ( @@ -159,7 +160,7 @@ def _store_type_embd_data(self) -> None: ) # Run through the strip embedding network embd_tensor = self.se_ttebd.embeddings_strip[0].call(two_side).detach() - torch.nn.Module.register_buffer(self, "type_embd_data", embd_tensor) + register_buffer_replacing_slot(self, "type_embd_data", embd_tensor) @cast_precision def call( diff --git a/deepmd/pt_expt/infer/deep_eval.py b/deepmd/pt_expt/infer/deep_eval.py index 9d6396d990..7cfa0d3c5b 100644 --- a/deepmd/pt_expt/infer/deep_eval.py +++ b/deepmd/pt_expt/infer/deep_eval.py @@ -691,27 +691,13 @@ def _load_pt(self, model_file: str, head: str | None = None) -> None: "sel": model.get_sel(), "dim_fparam": model.get_dim_fparam(), "dim_aparam": model.get_dim_aparam(), - "dim_chg_spin": model.get_dim_chg_spin() - if hasattr(model, "get_dim_chg_spin") - else 0, + "dim_chg_spin": model.get_dim_chg_spin(), "mixed_types": model.mixed_types(), "has_default_fparam": model.has_default_fparam(), "default_fparam": model.get_default_fparam(), - "has_chg_spin_ebd": ( - model.has_chg_spin_ebd() - if hasattr(model, "has_chg_spin_ebd") - else False - ), - "has_default_chg_spin": ( - model.has_default_chg_spin() - if hasattr(model, "has_default_chg_spin") - else False - ), - "default_chg_spin": ( - model.get_default_chg_spin() - if hasattr(model, "get_default_chg_spin") - else None - ), + "has_chg_spin_ebd": model.has_chg_spin_ebd(), + "has_default_chg_spin": model.get_default_chg_spin() is not None, + "default_chg_spin": model.get_default_chg_spin(), "is_spin": self._is_spin, "lower_input_kind": "graph" if use_graph_lower else "nlist", } @@ -909,14 +895,19 @@ def get_dim_aparam(self) -> int: def has_chg_spin_ebd(self) -> bool: """Check whether the model uses a dedicated charge_spin input.""" - if self._dpmodel is not None and hasattr(self._dpmodel, "has_chg_spin_ebd"): + if self._dpmodel is not None: return bool(self._dpmodel.has_chg_spin_ebd()) return bool(self.metadata.get("has_chg_spin_ebd", self.get_dim_chg_spin() > 0)) def has_default_chg_spin(self) -> bool: - """Check whether the model has a default charge_spin fallback.""" - if self._dpmodel is not None and hasattr(self._dpmodel, "has_default_chg_spin"): - return bool(self._dpmodel.has_default_chg_spin()) + """Check whether the model has a default charge_spin fallback. + + ``has_default_chg_spin`` was merged into ``get_default_chg_spin`` on + the live-model interfaces; this DeepEval wrapper method is kept for + API stability and computes the predicate directly. + """ + if self._dpmodel is not None: + return self._dpmodel.get_default_chg_spin() is not None return bool( self.metadata.get( "has_default_chg_spin", @@ -926,7 +917,7 @@ def has_default_chg_spin(self) -> bool: def get_dim_chg_spin(self) -> int: """Get the width of charge/spin condition inputs.""" - if self._dpmodel is not None and hasattr(self._dpmodel, "get_dim_chg_spin"): + if self._dpmodel is not None: return self._dpmodel.get_dim_chg_spin() return int(self.metadata.get("dim_chg_spin", 0)) @@ -967,6 +958,7 @@ def model_type(self) -> type["DeepEvalWrapper"]: """The evaluator of the model type.""" if self._dpmodel is not None: model_output_type = self._dpmodel.model_output_type() + var_name = self._dpmodel.get_var_name() else: # Metadata-only mode: derive the output-type set from the # fitting_output_defs names. `model_output_type()` on a @@ -975,6 +967,7 @@ def model_type(self) -> type["DeepEvalWrapper"]: model_output_type = [ d.name for d in self._model_output_def.def_outp.get_data().values() ] + var_name = None if "energy" in model_output_type: return DeepPot elif "dos" in model_output_type: @@ -985,11 +978,7 @@ def model_type(self) -> type["DeepEvalWrapper"]: return DeepPolar elif "wfc" in model_output_type: return DeepWFC - elif ( - self._dpmodel is not None - and hasattr(self._dpmodel, "get_var_name") - and self._dpmodel.get_var_name() in model_output_type - ): + elif var_name is not None and var_name in model_output_type: return DeepProperty else: raise RuntimeError("Unknown model type") @@ -1014,7 +1003,7 @@ def get_numb_dos(self) -> int: def get_var_name(self) -> str: """Get the name of the property (property models only).""" - if self._dpmodel is not None and hasattr(self._dpmodel, "get_var_name"): + if self._dpmodel is not None and self._dpmodel.get_var_name() is not None: return self._dpmodel.get_var_name() raise NotImplementedError( "get_var_name is only available for property models with the " @@ -1023,7 +1012,7 @@ def get_var_name(self) -> str: def get_task_dim(self) -> int: """Get the output dimension of the property (property models only).""" - if self._dpmodel is not None and hasattr(self._dpmodel, "get_task_dim"): + if self._dpmodel is not None: return self._dpmodel.get_task_dim() raise NotImplementedError( "get_task_dim is only available for property models with the " @@ -1032,7 +1021,7 @@ def get_task_dim(self) -> int: def get_intensive(self) -> bool: """Whether the property is intensive (property models only).""" - if self._dpmodel is not None and hasattr(self._dpmodel, "get_intensive"): + if self._dpmodel is not None: return self._dpmodel.get_intensive() raise NotImplementedError( "get_intensive is only available for property models with the " @@ -2419,7 +2408,7 @@ def _model_pair_excl(self) -> "PairExcludeMask | None": ) if self._dpmodel is not None: - pe = getattr(self._dpmodel.atomic_model, "pair_excl", None) + pe = self._dpmodel.atomic_model.pair_excl pet = pe.get_exclude_types() if pe is not None else [] else: pet = self.metadata.get("pair_exclude_types", []) @@ -2613,9 +2602,7 @@ def eval_descriptor( ext_atype_t, nlist_t, mapping=mapping_t, - charge_spin=charge_spin_t - if getattr(dp_am, "add_chg_spin_ebd", False) - else None, + charge_spin=charge_spin_t if dp_am.has_chg_spin_ebd() else None, ) return descriptor.detach().cpu().numpy() @@ -2684,9 +2671,7 @@ def eval_fitting_last_layer( ext_atype_t, nlist_t, mapping=mapping_t, - charge_spin=charge_spin_t - if getattr(dp_am, "add_chg_spin_ebd", False) - else None, + charge_spin=charge_spin_t if dp_am.has_chg_spin_ebd() else None, ) atype = ext_atype_t[:, :natoms] fitting_net = dp_am.fitting_net diff --git a/deepmd/pt_expt/model/make_model.py b/deepmd/pt_expt/model/make_model.py index c6def8f136..f59d507727 100644 --- a/deepmd/pt_expt/model/make_model.py +++ b/deepmd/pt_expt/model/make_model.py @@ -780,9 +780,10 @@ def _call_common_graph( with_csr = ( not self.training and cuda_infer_level() >= 1 - and bool(getattr(_desc, "geo_compress", False)) + and _desc is not None + and _desc.get_geo_compress() ) - pair_excl = getattr(self.atomic_model, "pair_excl", None) + pair_excl = self.atomic_model.pair_excl ng = build_neighbor_graph_for_method( method, cc, atype, bb, rcut, pair_excl, with_csr=with_csr ) diff --git a/deepmd/pt_expt/train/training.py b/deepmd/pt_expt/train/training.py index f277452fcb..c9e73f03c1 100644 --- a/deepmd/pt_expt/train/training.py +++ b/deepmd/pt_expt/train/training.py @@ -371,9 +371,9 @@ def get_additional_data_requirement(_model: Any) -> list[DataRequirementItem]: ) ) if _model.has_chg_spin_ebd(): - has_default_cs = _model.has_default_chg_spin() + default_cs = _model.get_default_chg_spin() + has_default_cs = default_cs is not None if has_default_cs: - default_cs = _model.get_default_chg_spin() if hasattr(default_cs, "cpu"): default_cs = default_cs.cpu().numpy() else: @@ -1088,7 +1088,7 @@ def forward( distinguish_types=False, # model-level pair exclusion is a nlist-BUILD transform (decision # #18/A4); the compiled dense lower consumes a pre-excluded nlist. - pair_excl=getattr(self.original_model.atomic_model, "pair_excl", None), + pair_excl=self.original_model.atomic_model.pair_excl, ) ext_coord = ext_coord.reshape(nframes, -1, 3) @@ -1125,9 +1125,7 @@ def forward( .reshape(1, _dim_fparam) .expand(nframes, -1) ) - _dim_cs = ( - _model.get_dim_chg_spin() if hasattr(_model, "get_dim_chg_spin") else 0 - ) + _dim_cs = _model.get_dim_chg_spin() if charge_spin is None and _dim_cs > 0: _default_cs = _model.get_default_chg_spin() if _default_cs is not None: @@ -1294,9 +1292,7 @@ def _forward_graph( .reshape(1, _dim_fparam) .expand(nframes, -1) ) - _dim_cs = ( - _model.get_dim_chg_spin() if hasattr(_model, "get_dim_chg_spin") else 0 - ) + _dim_cs = _model.get_dim_chg_spin() if charge_spin is None and _dim_cs > 0: _default_cs = _model.get_default_chg_spin() if _default_cs is not None: @@ -1313,7 +1309,7 @@ def _forward_graph( # level pair_exclude is a graph-BUILD transform (decision #18): fold it # into edge_mask here so the compiled lower consumes a pre-excluded graph # (the lower no longer re-applies it), matching the eager path exactly. - pair_excl = getattr(_model.atomic_model, "pair_excl", None) + pair_excl = _model.atomic_model.pair_excl ng = build_neighbor_graph_for_method( getattr(_model, "neighbor_graph_method", "dense"), coord_3d, diff --git a/deepmd/pt_expt/utils/network.py b/deepmd/pt_expt/utils/network.py index 004ba94401..f3ca3b392d 100644 --- a/deepmd/pt_expt/utils/network.py +++ b/deepmd/pt_expt/utils/network.py @@ -99,7 +99,7 @@ def __setattr__(self, name: str, value: Any) -> None: self._buffers[name] = None return None return super().__setattr__(name, None) - if getattr(self, "trainable", False): + if self.trainable: param = ( value if isinstance(value, torch.nn.Parameter) diff --git a/deepmd/pt_expt/utils/serialization.py b/deepmd/pt_expt/utils/serialization.py index e82ae24e0a..5d7be558bb 100644 --- a/deepmd/pt_expt/utils/serialization.py +++ b/deepmd/pt_expt/utils/serialization.py @@ -394,7 +394,7 @@ def _make_sample_inputs( else: aparam = None - dim_chg_spin = model.get_dim_chg_spin() if hasattr(model, "get_dim_chg_spin") else 0 + dim_chg_spin = model.get_dim_chg_spin() if dim_chg_spin > 0: charge_spin = torch.zeros( nframes, dim_chg_spin, dtype=torch.float64, device=_env.DEVICE @@ -515,7 +515,7 @@ def build_synthetic_graph_inputs( ntypes = len(model.get_type_map()) dim_fparam = model.get_dim_fparam() dim_aparam = model.get_dim_aparam() - dim_chg_spin = model.get_dim_chg_spin() if hasattr(model, "get_dim_chg_spin") else 0 + dim_chg_spin = model.get_dim_chg_spin() # Box large enough to avoid PBC degeneracy; centered coords. box_size = rcut * 3.0 @@ -1054,25 +1054,13 @@ def _collect_metadata( "nnei": sum(model.get_sel()), "dim_fparam": model.get_dim_fparam(), "dim_aparam": model.get_dim_aparam(), - "dim_chg_spin": ( - model.get_dim_chg_spin() if hasattr(model, "get_dim_chg_spin") else 0 - ), + "dim_chg_spin": model.get_dim_chg_spin(), "mixed_types": model.mixed_types(), "has_default_fparam": model.has_default_fparam(), "default_fparam": model.get_default_fparam(), - "has_chg_spin_ebd": ( - model.has_chg_spin_ebd() if hasattr(model, "has_chg_spin_ebd") else False - ), - "has_default_chg_spin": ( - model.has_default_chg_spin() - if hasattr(model, "has_default_chg_spin") - else False - ), - "default_chg_spin": ( - _metadata_value_to_json(model.get_default_chg_spin()) - if hasattr(model, "get_default_chg_spin") - else None - ), + "has_chg_spin_ebd": model.has_chg_spin_ebd(), + "has_default_chg_spin": model.get_default_chg_spin() is not None, + "default_chg_spin": _metadata_value_to_json(model.get_default_chg_spin()), "fitting_output_defs": fitting_output_defs, # sel_type enables `DeepEval.get_sel_type()` without a dpmodel # round-trip; required for dipole/polar/wfc models in metadata-only @@ -1101,15 +1089,21 @@ def _collect_metadata( # ``atomic_model.has_message_passing()`` is important for composite # atomic models (e.g. ``LinearAtomicModel`` in DP-ZBL) which don't # expose a single ``.descriptor`` but do aggregate the flag across - # their sub-models. ``descriptor.has_message_passing()`` is the - # fallback for any future wrapper that lacks the higher-level - # methods. + # their sub-models. ``has_message_passing`` is declared on the base + # model/atomic-model/descriptor classes, so every concrete object at + # each level implements it; ``descriptor.has_message_passing()`` only + # matters as a fallback when an upstream level raises + # ``NotImplementedError`` (e.g. an atomic model that intentionally + # opts out), never for a missing method. def _probe_has_message_passing(obj: object) -> bool | None: - if obj is None or not hasattr(obj, "has_message_passing"): + # has_message_passing is @abstractmethod on the base descriptor, so + # every concrete descriptor implements it; a wrapper lacking it is a + # construction bug that must raise, not degrade silently. + if obj is None: return None try: return bool(obj.has_message_passing()) - except (AttributeError, NotImplementedError): + except NotImplementedError: return None result: bool | None = None @@ -1140,12 +1134,19 @@ def _probe_has_message_passing(obj: object) -> bool | None: # feeders (C++ ``DeepPotPTExpt::init``, metadata-only DeepEval) rebuild # the mask. Descriptor-level ``exclude_types`` needs NO metadata: it is # fully inside the compiled artifact. + from deepmd.dpmodel.atomic_model.base_atomic_model import ( + BaseAtomicModel, + ) + pair_exclude_types: list[list[int]] = [] for obj in ( getattr(model, "atomic_model", None), model, ): - pet = getattr(obj, "pair_exclude_types", None) + # `obj` may be the atomic model (the owner of pair_exclude_types) or + # the full model (e.g. the ``model`` fallback above); only the former + # implements the accessor, so gate on it instead of getattr-probing. + pet = obj.get_pair_exclude_types() if isinstance(obj, BaseAtomicModel) else None if pet: pair_exclude_types = [[int(ti), int(tj)] for (ti, tj) in pet] break @@ -1542,7 +1543,7 @@ def _trace_and_export( ) _forbidden = forbidden_dims_from_model(model) - _dim_cs = model.get_dim_chg_spin() if hasattr(model, "get_dim_chg_spin") else 0 + _dim_cs = model.get_dim_chg_spin() if _dim_cs > 1: _forbidden.add(int(_dim_cs)) nframes_sample = next_safe_prime(5, _forbidden) diff --git a/deepmd/tf2/model/dp_model.py b/deepmd/tf2/model/dp_model.py index 1326e1c881..bd69555dc4 100644 --- a/deepmd/tf2/model/dp_model.py +++ b/deepmd/tf2/model/dp_model.py @@ -87,6 +87,8 @@ def call_common( # ``_input_type_cast`` (dpmodel make_model) returns a ``spin`` slot # for the native-spin graph route; tf2 has no spin/graph lower, so # it is discarded here. + # Guard atomic_model too: test doubles may lack it. + am = getattr(self, "atomic_model", None) cc, bb, fp, ap, cs, _, input_prec = self._input_type_cast( to_tensorflow_array(coord), box=to_tensorflow_array(box), @@ -113,11 +115,8 @@ def call_common( # Model-level pair exclusion is a nlist-BUILD transform # (decision #18/A4): fold it into the freshly built nlist here so # the live/eager TF2 upper path matches the SavedModel export and - # the other backends. Identity when nothing is excluded. Guard - # atomic_model too: test doubles may lack it. - pair_excl=getattr( - getattr(self, "atomic_model", None), "pair_excl", None - ), + # the other backends. Identity when nothing is excluded. + pair_excl=am.pair_excl if am is not None else None, pass_lower_kwargs=True, ) return self._output_type_cast(model_predict, input_prec) diff --git a/deepmd/tf2/train/trainer.py b/deepmd/tf2/train/trainer.py index 4c1e8af671..81a1139221 100644 --- a/deepmd/tf2/train/trainer.py +++ b/deepmd/tf2/train/trainer.py @@ -164,9 +164,10 @@ def get_additional_data_requirement(_model: Any) -> list[DataRequirementItem]: ) ) if _model.has_chg_spin_ebd(): - has_default_cs = _model.has_default_chg_spin() + default_chg_spin = _model.get_default_chg_spin() + has_default_cs = default_chg_spin is not None default_cs = ( - np.asarray(to_tf_tensor(_model.get_default_chg_spin()).numpy()) + np.asarray(to_tf_tensor(default_chg_spin).numpy()) if has_default_cs else 0.0 ) @@ -882,6 +883,8 @@ def compiled_prepare_lower_batch( aparam=to_tensorflow_array(aparam), charge_spin=to_tensorflow_array(charge_spin), ) + # Guard atomic_model for test doubles. + am = getattr(model, "atomic_model", None) return prepare_lower_inputs( rcut=model.get_rcut(), sel=model.get_sel(), @@ -895,10 +898,8 @@ def compiled_prepare_lower_batch( # Model-level pair exclusion is a nlist-BUILD transform # (decision #18/A4): the compiled lower consumes a pre-excluded # nlist, so fold exclusion in here at the compiled-training - # prepare seam. Guard atomic_model for test doubles. - pair_excl=getattr( - getattr(model, "atomic_model", None), "pair_excl", None - ), + # prepare seam. + pair_excl=am.pair_excl if am is not None else None, ) return compiled_prepare_lower_batch diff --git a/deepmd/tf2/utils/serialization.py b/deepmd/tf2/utils/serialization.py index baedd5ba87..63bbdf6583 100644 --- a/deepmd/tf2/utils/serialization.py +++ b/deepmd/tf2/utils/serialization.py @@ -377,6 +377,10 @@ def call( fparam: tf.Tensor, aparam: tf.Tensor, ) -> dict[str, tf.Tensor]: + # exclusion is a nlist-BUILD transform (decision #18/A4); the + # traced lower consumes a pre-excluded nlist. Guard atomic_model + # too: test doubles (DummyModel) lack it. + am = getattr(model, "atomic_model", None) return unwrap_value( model_call_from_call_lower( call_lower=call_lower, @@ -390,12 +394,7 @@ def call( fparam=fparam, aparam=aparam, do_atomic_virial=do_atomic_virial, - # exclusion is a nlist-BUILD transform (decision #18/A4); - # the traced lower consumes a pre-excluded nlist. Guard - # atomic_model too: test doubles (DummyModel) lack it. - pair_excl=getattr( - getattr(model, "atomic_model", None), "pair_excl", None - ), + pair_excl=am.pair_excl if am is not None else None, ) ) @@ -517,7 +516,8 @@ def get_pair_exclude_types() -> tf.Tensor: # traced call_lower_* consumes it (decision #18/A4). The compiled ``call`` # already pre-excludes its freshly built nlist. Guard atomic_model: test # doubles may lack it. - pet = getattr(getattr(model, "atomic_model", None), "pair_exclude_types", []) + am = getattr(model, "atomic_model", None) + pet = am.get_pair_exclude_types() if am is not None else [] flat = [int(t) for pair in (pet or []) for t in pair] return tf.constant(flat, dtype=tf.int64) @@ -555,7 +555,12 @@ def get_default_fparam() -> tf.Tensor: # property models: persist the output name/dimension/intensiveness so the # evaluator can dispatch to DeepProperty and reshape the output. - if hasattr(model, "get_var_name"): + # ``get_var_name`` is declared on every model with a concrete default + # of ``None`` for non-property models (issue #5897), so a bare + # ``hasattr`` check is always true; gate on the actual return value + # instead, mirroring deepmd/jax/jax2tf/serialization.py. + is_property_model = model.get_var_name() is not None + if is_property_model: @tf.function def get_var_name() -> tf.Tensor: diff --git a/source/jax2tf_tests/test_serialization.py b/source/jax2tf_tests/test_serialization.py index 27751b7c43..bc4cf02239 100644 --- a/source/jax2tf_tests/test_serialization.py +++ b/source/jax2tf_tests/test_serialization.py @@ -28,6 +28,22 @@ def _saved_model_ops(model_dir: Path) -> set[str]: return ops +def test_decode_default_chg_spin_preserves_the_no_default_marker() -> None: + """The exporter writes an EMPTY tensor when the model has no default + charge-spin; the wrapper must decode it back to ``None`` so the + live-model invariant ``get_default_chg_spin() is None == no default`` + holds at the artifact boundary too. + """ + pytest.importorskip("jax") + + from deepmd.jax.jax2tf.tfmodel import ( + _decode_default_chg_spin, + ) + + assert _decode_default_chg_spin(False, []) is None + assert _decode_default_chg_spin(True, [2.0, 1.0]) == [2.0, 1.0] + + def test_savedmodel_export_contains_xla_call_module(tmp_path, monkeypatch) -> None: pytest.importorskip("jax") pytest.importorskip("flax") @@ -122,6 +138,10 @@ def has_default_chg_spin(self) -> bool: def get_default_chg_spin(self) -> None: return None + def get_var_name(self) -> None: + # non-property model, matching the make_base_model default + return None + class DummyChargeSpinModel(DummyModel): dim_chg_spin = 2 diff --git a/source/tests/common/dpmodel/test_descriptor_block_defaults.py b/source/tests/common/dpmodel/test_descriptor_block_defaults.py new file mode 100644 index 0000000000..daef3c46aa --- /dev/null +++ b/source/tests/common/dpmodel/test_descriptor_block_defaults.py @@ -0,0 +1,94 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Class-level defaults for the ``set_davg_zero`` / ``set_stddev_constant`` +stat-behavior flags on ``DescriptorBlock`` (issue #5897): stat machinery +must be able to read these flags on any block without a ``getattr`` probe. +""" + +import numpy as np + +from deepmd.dpmodel.descriptor import ( + DescrptSeA, +) +from deepmd.dpmodel.descriptor.descriptor import ( + DescriptorBlock, +) +from deepmd.dpmodel.descriptor.make_base_descriptor import ( + make_base_descriptor, +) +from deepmd.dpmodel.utils.env_mat_stat import ( + merge_env_stat, +) + + +def test_block_stat_flags_have_class_defaults() -> None: + """Every DescriptorBlock answers the stat-behavior flags without a + getattr probe: class-level defaults False, blocks override in __init__. + """ + assert DescriptorBlock.set_davg_zero is False + assert DescriptorBlock.set_stddev_constant is False + + +def test_base_descriptor_stat_flags_have_class_defaults() -> None: + """The ``BD`` base in ``make_base_descriptor`` (the ``Descriptor``-side + twin of ``DescriptorBlock`` above) carries the same concrete class + defaults, so ``merge_env_stat`` -- which accepts either a ``Descriptor`` + or a ``DescriptorBlock`` as ``base_obj`` -- can read the flags on a bare + ``Descriptor`` without a ``getattr`` probe. + """ + bd = make_base_descriptor(np.ndarray, "call") + assert bd.set_davg_zero is False + assert bd.set_stddev_constant is False + + +def _sample() -> dict: + rng = np.random.default_rng(0) + nf, nloc = 2, 6 + coord = rng.normal(size=(nf, nloc, 3)) * 2.0 + atype = np.array([[0, 1, 0, 1, 0, 1], [1, 0, 1, 0, 1, 0]], dtype=np.int64) + box = np.tile((np.eye(3) * 12.0).reshape(1, 9), (nf, 1)) + return {"coord": coord, "atype": atype, "box": box} + + +def test_merge_env_stat_on_bare_descriptor_no_attribute_error() -> None: + """``merge_env_stat`` reads ``base_obj.set_davg_zero`` / + ``set_stddev_constant`` unconditionally (no ``getattr`` probe). Pin that + this does not raise ``AttributeError`` when ``base_obj`` is a bare + se-family ``Descriptor`` (not a ``DescriptorBlock``) which never sets + those flags itself and instead relies on the ``BD`` base's class + defaults. + """ + base = DescrptSeA(6.0, 0.5, [10, 10]) + link = DescrptSeA(6.0, 0.5, [10, 10]) + sample = _sample() + base.compute_input_stats([sample]) + link.compute_input_stats([sample]) + # Would raise AttributeError before the BD-base class defaults existed + # if a descriptor never assigned instance attributes for these flags. + merge_env_stat(base, link) + + +def test_block_stat_flags_override_branch() -> None: + """A block constructed with ``set_davg_zero=True`` shadows the class + default with an instance attribute; a block constructed with the + default arguments keeps reading the class default (False). + """ + from deepmd.dpmodel.descriptor.dpa1 import ( + DescrptBlockSeAtten, + ) + + blk_default = DescrptBlockSeAtten( + rcut=4.0, + rcut_smth=0.5, + sel=[6, 6], + ntypes=2, + ) + assert blk_default.set_davg_zero is False + + blk_override = DescrptBlockSeAtten( + rcut=4.0, + rcut_smth=0.5, + sel=[6, 6], + ntypes=2, + set_davg_zero=True, + ) + assert blk_override.set_davg_zero is True diff --git a/source/tests/common/dpmodel/test_descriptor_dpa1.py b/source/tests/common/dpmodel/test_descriptor_dpa1.py index 1df473c087..371bd69c85 100644 --- a/source/tests/common/dpmodel/test_descriptor_dpa1.py +++ b/source/tests/common/dpmodel/test_descriptor_dpa1.py @@ -5,6 +5,7 @@ from deepmd.dpmodel.descriptor import ( DescrptDPA1, + DescrptSeA, ) from ...seed import ( @@ -65,6 +66,27 @@ def test_lmax_two_serialization(self) -> None: for index in (0, 1, 4): np.testing.assert_allclose(actual[index], expected[index]) + def test_tebd_compression_slots_declared(self) -> None: + """Tebd-compression slots are class properties, not runtime accidents. + + ``type_embd_data``/``tebd_compress`` must be declared (defaulted) + in ``__init__`` of tebd-family descriptors so their presence does + not depend on whether compression was ever enabled. The jax + restore walker (``deepmd/jax/utils/serialization.py``) relies on + ``hasattr(obj, "tebd_compress")`` as a family-membership test, so + a non-tebd descriptor (e.g. ``DescrptSeA``) must never carry + either attribute (see issue #5897). + """ + em0 = DescrptDPA1(self.rcut, self.rcut_smth, self.sel, ntypes=2) + self.assertIsNone(em0.type_embd_data) + self.assertFalse(em0.tebd_compress) + self.assertIsNone(em0.se_atten.type_embd_data) + self.assertFalse(em0.se_atten.tebd_compress) + + se_a = DescrptSeA(self.rcut, self.rcut_smth, self.sel) + self.assertFalse(hasattr(se_a, "type_embd_data")) + self.assertFalse(hasattr(se_a, "tebd_compress")) + def test_multiple_frames(self) -> None: rng = np.random.default_rng(GLOBAL_SEED) nf, nloc, nnei = self.nlist.shape diff --git a/source/tests/common/dpmodel/test_make_base_fitting.py b/source/tests/common/dpmodel/test_make_base_fitting.py new file mode 100644 index 0000000000..0eacd2b117 --- /dev/null +++ b/source/tests/common/dpmodel/test_make_base_fitting.py @@ -0,0 +1,39 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +import numpy as np +import pytest + +from deepmd.dpmodel.fitting.make_base_fitting import ( + make_base_fitting, +) + + +class MinimalFitting(make_base_fitting(np.ndarray)): + """Smallest concrete fitting: exercises the base-declared defaults.""" + + def output_def(self): + raise NotImplementedError + + def fwd(self, *args, **kwargs): + raise NotImplementedError + + def get_type_map(self): + return [] + + def change_type_map(self, type_map, model_with_new_type_stat=None): + raise NotImplementedError + + def serialize(self): + return {} + + @classmethod + def deserialize(cls, data): + return cls() + + +def test_reinit_exclude_default_noop_on_empty() -> None: + MinimalFitting().reinit_exclude([]) + + +def test_reinit_exclude_default_raises_on_nonempty() -> None: + with pytest.raises(NotImplementedError): + MinimalFitting().reinit_exclude([0]) diff --git a/source/tests/common/dpmodel/test_model_compression.py b/source/tests/common/dpmodel/test_model_compression.py index 29f5c281a8..8599d3e51b 100644 --- a/source/tests/common/dpmodel/test_model_compression.py +++ b/source/tests/common/dpmodel/test_model_compression.py @@ -344,6 +344,8 @@ def test_se_atten_enable_compression(self) -> None: self.assertTrue(compressed.compress) self.assertTrue(compressed.geo_compress) + self.assertTrue(compressed.get_geo_compress()) + self.assertFalse(descriptor.get_geo_compress()) serialized = compressed.serialize() self.assertEqual(serialized["@version"], 3) self.assertIn("compress", serialized) diff --git a/source/tests/common/dpmodel/test_pair_exclude_contract.py b/source/tests/common/dpmodel/test_pair_exclude_contract.py new file mode 100644 index 0000000000..ce2e21bdde --- /dev/null +++ b/source/tests/common/dpmodel/test_pair_exclude_contract.py @@ -0,0 +1,60 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Pin the ``pair_exclude_types``/``pair_excl`` construction-path contract +(issue #5897 task 5): ``BaseAtomicModel.__init__`` (via +``reinit_pair_exclude``) always sets both attributes, on every construction +path (direct ``__init__`` and ``deserialize``); ``get_pair_exclude_types()`` +is the public accessor, ``pair_excl`` a pinned direct-access attribute. +""" + +from deepmd.dpmodel.atomic_model import ( + DPAtomicModel, +) +from deepmd.dpmodel.descriptor import ( + DescrptSeA, +) +from deepmd.dpmodel.fitting import ( + InvarFitting, +) + +RCUT = 2.2 +RCUT_SMTH = 0.4 +SEL = [5, 2] +NTYPES = 2 +TYPE_MAP = ["foo", "bar"] + + +def _make_minimal_atomic_model( + pair_exclude_types: list[tuple[int, int]], +) -> DPAtomicModel: + ds = DescrptSeA( + RCUT, + RCUT_SMTH, + SEL, + ) + ft = InvarFitting( + "energy", + NTYPES, + ds.get_dim_out(), + 1, + mixed_types=ds.mixed_types(), + ) + return DPAtomicModel( + ds, + ft, + type_map=TYPE_MAP, + pair_exclude_types=pair_exclude_types, + ) + + +def test_pair_excl_exists_after_init_and_deserialize() -> None: + md0 = _make_minimal_atomic_model(pair_exclude_types=[(0, 1)]) + assert md0.get_pair_exclude_types() == [(0, 1)] + assert md0.pair_excl is not None + # the mask holds the SYMMETRIC CLOSURE of the accessor's pairs + assert {tuple(p) for p in md0.pair_excl.get_exclude_types()} == {(0, 1), (1, 0)} + md1 = type(md0).deserialize(md0.serialize()) + assert md1.get_pair_exclude_types() == [(0, 1)] + assert md1.pair_excl is not None + md2 = _make_minimal_atomic_model(pair_exclude_types=[]) + assert md2.get_pair_exclude_types() == [] + assert md2.pair_excl is None diff --git a/source/tests/common/dpmodel/test_zbl_bridging.py b/source/tests/common/dpmodel/test_zbl_bridging.py index 0db497dd84..cf810a4cb9 100644 --- a/source/tests/common/dpmodel/test_zbl_bridging.py +++ b/source/tests/common/dpmodel/test_zbl_bridging.py @@ -396,7 +396,7 @@ def test_default_conditioning_accessors_are_forwarded(self) -> None: """``has_default_*`` must not fall through to the base either.""" bridged = self._model(bridging=True) plain = self._model(bridging=False) - assert bridged.has_default_chg_spin() == plain.has_default_chg_spin() + assert bridged.get_default_chg_spin() == plain.get_default_chg_spin() assert bridged.has_default_fparam() == plain.has_default_fparam() assert bridged.get_default_fparam() == plain.get_default_fparam() diff --git a/source/tests/infer/gen_dpa4_spin_chgspin.py b/source/tests/infer/gen_dpa4_spin_chgspin.py index 502767c6ac..e8aa53a359 100644 --- a/source/tests/infer/gen_dpa4_spin_chgspin.py +++ b/source/tests/infer/gen_dpa4_spin_chgspin.py @@ -125,7 +125,7 @@ def _build_model_dict() -> dict: f"expected the combined native-spin DPA4 to expose dim_chg_spin == 2, " f"got {model.get_dim_chg_spin()}" ) - assert model.has_default_chg_spin() + assert model.get_default_chg_spin() is not None model_dict = model.serialize() model_dict = jitter_zero_arrays(model_dict, np.random.default_rng(_JITTER_SEED)) return model_dict diff --git a/source/tests/jax/test_training.py b/source/tests/jax/test_training.py index 0055a35b4b..913a851fe3 100644 --- a/source/tests/jax/test_training.py +++ b/source/tests/jax/test_training.py @@ -496,6 +496,11 @@ def __call__(self) -> tuple[np.ndarray, np.ndarray]: class _DescriptorWithStats: + # stat-behavior flags merge_env_stat reads directly on any + # Descriptor/DescriptorBlock (class defaults on the real bases) + set_davg_zero = False + set_stddev_constant = False + def __init__(self, stats: dict[str, StatItem]) -> None: self.stats = stats self.davg = np.asarray([0.0], dtype=np.float64) diff --git a/source/tests/pt_expt/model/test_dpa4_native_spin.py b/source/tests/pt_expt/model/test_dpa4_native_spin.py index e16a08172a..d88ed55617 100644 --- a/source/tests/pt_expt/model/test_dpa4_native_spin.py +++ b/source/tests/pt_expt/model/test_dpa4_native_spin.py @@ -1134,7 +1134,7 @@ def test_training_smoke_combined(self, tmp_path) -> None: model = trainer.wrapper.model[DEFAULT_TASK_KEY] assert isinstance(model, NativeSpinEnergyModel) assert model.has_chg_spin_ebd() - assert model.has_default_chg_spin() + assert model.get_default_chg_spin() is not None tasks = trainer._make_training_tasks() task = trainer.select_task(tasks) diff --git a/source/tests/universal/common/cases/atomic_model/utils.py b/source/tests/universal/common/cases/atomic_model/utils.py index 45b30c6454..fbf6600320 100644 --- a/source/tests/universal/common/cases/atomic_model/utils.py +++ b/source/tests/universal/common/cases/atomic_model/utils.py @@ -97,6 +97,26 @@ def test_has_message_passing(self) -> None: module.has_message_passing(), self.expected_has_message_passing ) + def test_pair_exclude_contract(self) -> None: + """``pair_exclude_types``/``pair_excl`` are set by + ``BaseAtomicModel.__init__`` on every construction path; the + accessor is the public surface, the mask attribute is internal. + """ + pet = self.module.get_pair_exclude_types() + assert isinstance(pet, list) + # invariant pinned by reinit_pair_exclude: + assert (self.module.pair_excl is None) == (len(pet) == 0) + if self.module.pair_excl is not None: + # the mask must hold exactly the SYMMETRIC CLOSURE of the + # accessor's pairs (both backends symmetrize on reinit; + # compare as sets of tuples: pt stores a set, dpmodel a list) + symmetrized = { + pair for i, j in map(tuple, pet) for pair in ((i, j), (j, i)) + } + assert { + tuple(p) for p in self.module.pair_excl.get_exclude_types() + } == symmetrized + def test_forward(self) -> None: """Test forward.""" nf = 1 diff --git a/source/tests/universal/common/cases/descriptor/utils.py b/source/tests/universal/common/cases/descriptor/utils.py index 03dd5248da..ca5a750e1c 100644 --- a/source/tests/universal/common/cases/descriptor/utils.py +++ b/source/tests/universal/common/cases/descriptor/utils.py @@ -48,6 +48,20 @@ def test_capability_contract(self) -> None: # off -- a no-op on descriptors without one. self.module.disable_graph_lower() assert self.module.uses_graph_lower() is False + # chg-spin family: concrete base defaults, never probed (issue #5897). + assert isinstance(self.module.get_dim_chg_spin(), int) + dcs = self.module.get_default_chg_spin() + # Backend-agnostic: dpmodel returns a list, frozen pt returns a + # torch.Tensor -- pin the shape contract, not the container type. + assert dcs is None or len(dcs) == self.module.get_dim_chg_spin() + # has_default_chg_spin was merged into get_default_chg_spin: the + # predicate is ``get_default_chg_spin() is not None``. The absence + # of ``has_default_chg_spin`` on the dpmodel side is pinned in + # source/tests/universal/dpmodel/descriptor/test_descriptor.py -- + # pt is frozen and still declares the (now-redundant) method. + # Geometric-compression state query: base default False, the + # dpa1/dpa2 families override from their ``geo_compress`` attribute. + assert isinstance(self.module.get_geo_compress(), bool) def test_forward_consistency(self) -> None: ret = [] diff --git a/source/tests/universal/common/cases/fitting/utils.py b/source/tests/universal/common/cases/fitting/utils.py index de6b12c3a2..eae1cdde3c 100644 --- a/source/tests/universal/common/cases/fitting/utils.py +++ b/source/tests/universal/common/cases/fitting/utils.py @@ -104,6 +104,17 @@ def test_exclude_types( )[var_name] np.testing.assert_allclose(rd, rd_ex) + def test_reinit_exclude_contract(self) -> None: + """``reinit_exclude`` is declared on the base fitting: empty input + is always accepted; fittings with exclusion support apply it. + """ + self.module.reinit_exclude([]) # must never raise + if hasattr(self.module, "emask"): # override branch + self.module.reinit_exclude([0]) + assert self.module.exclude_types == [0] + self.module.reinit_exclude([]) + assert self.module.exclude_types == [] + def test_change_type_map(self) -> None: if not self.module.mixed_types: # skip if not mixed_types diff --git a/source/tests/universal/common/cases/model/utils.py b/source/tests/universal/common/cases/model/utils.py index 5ff4254934..c51f9aefd1 100644 --- a/source/tests/universal/common/cases/model/utils.py +++ b/source/tests/universal/common/cases/model/utils.py @@ -131,6 +131,38 @@ def test_has_spin(self) -> None: expected = getattr(self, "test_spin", False) self.assertEqual(self.module.has_spin(), expected) + def test_chg_spin_capability_contract(self) -> None: + """chg-spin queries are declared on the base model with concrete + defaults (False/0/None) -- direct calls, never ``hasattr`` probes. + + ``has_default_chg_spin`` was merged into ``get_default_chg_spin`` + (predicate: ``get_default_chg_spin() is not None``). Its absence on + the dpmodel side is pinned in + source/tests/universal/dpmodel/model/test_model.py -- pt is frozen + and still declares the (now-redundant) method. + """ + assert isinstance(self.module.has_chg_spin_ebd(), bool) + assert isinstance(self.module.get_dim_chg_spin(), int) + dcs = self.module.get_default_chg_spin() + # Backend-agnostic: dpmodel returns a list, frozen pt returns a + # torch.Tensor -- pin the shape contract, not the container type. + assert dcs is None or len(dcs) == self.module.get_dim_chg_spin() + + def test_property_capability_contract(self) -> None: + """Property queries are declared on the base model with concrete + defaults: ``get_var_name`` returns ``None`` for non-property models + (the support predicate), ``get_intensive`` defaults ``False``, + ``get_task_dim`` raises for non-property models. + """ + vn = self.module.get_var_name() + assert vn is None or isinstance(vn, str) + assert isinstance(self.module.get_intensive(), bool) + if vn is None: + with self.assertRaises(NotImplementedError): + self.module.get_task_dim() + else: + assert isinstance(self.module.get_task_dim(), int) + def test_forward(self) -> None: """Test forward and forward_lower.""" test_spin = getattr(self, "test_spin", False) diff --git a/source/tests/universal/dpmodel/descriptor/test_descriptor.py b/source/tests/universal/dpmodel/descriptor/test_descriptor.py index 6999c0b779..85faf46ed9 100644 --- a/source/tests/universal/dpmodel/descriptor/test_descriptor.py +++ b/source/tests/universal/dpmodel/descriptor/test_descriptor.py @@ -19,6 +19,9 @@ DescrptSeT, DescrptSeTTebd, ) +from deepmd.dpmodel.descriptor.base_descriptor import ( + BaseDescriptor, +) from deepmd.dpmodel.descriptor.dpa2 import ( RepformerArgs, RepinitArgs, @@ -26,6 +29,9 @@ from deepmd.dpmodel.descriptor.dpa3 import ( RepFlowArgs, ) +from deepmd.dpmodel.descriptor.make_base_descriptor import ( + make_base_descriptor, +) from deepmd.dpmodel.descriptor.repflows import ( DescrptBlockRepflows, ) @@ -999,17 +1005,47 @@ def test_shared_default_required_for_hybrid_default(self) -> None: shared_default = DescrptHybrid( list=[self._make_dpa3([5.0, 1.0]), self._make_dpa3([5.0, 1.0])] ) - self.assertTrue(shared_default.has_default_chg_spin()) + self.assertIsNotNone(shared_default.get_default_chg_spin()) self.assertEqual(shared_default.get_default_chg_spin(), [5.0, 1.0]) missing_default = DescrptHybrid( list=[self._make_dpa3([5.0, 1.0]), self._make_dpa3(None)] ) - self.assertFalse(missing_default.has_default_chg_spin()) self.assertIsNone(missing_default.get_default_chg_spin()) mismatched_default = DescrptHybrid( list=[self._make_dpa3([5.0, 1.0]), self._make_dpa3([6.0, 1.0])] ) - self.assertFalse(mismatched_default.has_default_chg_spin()) self.assertIsNone(mismatched_default.get_default_chg_spin()) + + +class TestHasDefaultChgSpinAbsentDP(unittest.TestCase): + """Pin the dpmodel-side half of the ``has_default_chg_spin`` merge. + + ``has_default_chg_spin`` was merged into ``get_default_chg_spin`` + (issue #5897): the predicate is ``get_default_chg_spin() is not None``. + The shared universal descriptor case only asserts the concrete + replacement (``get_default_chg_spin``), since the frozen pt backend + still declares the (now-redundant) ``has_default_chg_spin`` method on + several descriptors. This test pins that the method is gone from the + dpmodel base-descriptor family, where the merge is authoritative. + """ + + def test_absent_from_base_descriptor(self) -> None: + assert not hasattr(BaseDescriptor, "has_default_chg_spin") + assert not hasattr( + make_base_descriptor(np.ndarray, "call"), "has_default_chg_spin" + ) + + def test_absent_from_concrete_descriptors(self) -> None: + for cls in ( + DescrptSeA, + DescrptSeR, + DescrptSeT, + DescrptSeTTebd, + DescrptDPA1, + DescrptDPA2, + DescrptDPA3, + DescrptHybrid, + ): + assert not hasattr(cls, "has_default_chg_spin") diff --git a/source/tests/universal/dpmodel/model/test_model.py b/source/tests/universal/dpmodel/model/test_model.py index ece94a58ad..a51611069d 100644 --- a/source/tests/universal/dpmodel/model/test_model.py +++ b/source/tests/universal/dpmodel/model/test_model.py @@ -18,6 +18,10 @@ EnergyModel, SpinModel, ) +from deepmd.dpmodel.model.base_model import ( + BaseModel, + make_base_model, +) from deepmd.utils.spin import ( Spin, ) @@ -271,3 +275,24 @@ def setUpClass(cls) -> None: cls.expected_dim_fparam = ft.get_dim_fparam() cls.expected_dim_aparam = ft.get_dim_aparam() cls.skip_test_autodiff = True + + +class TestHasDefaultChgSpinAbsentDP(unittest.TestCase): + """Pin the dpmodel-side half of the ``has_default_chg_spin`` merge. + + ``has_default_chg_spin`` was merged into ``get_default_chg_spin`` + (issue #5897): the predicate is ``get_default_chg_spin() is not None``. + The shared universal model case only asserts the concrete replacement + (``get_default_chg_spin``), since the frozen pt backend still declares + the (now-redundant) ``has_default_chg_spin`` method on ``make_model``. + This test pins that the method is gone from the dpmodel base-model + family, where the merge is authoritative. + """ + + def test_absent_from_base_model(self) -> None: + assert not hasattr(BaseModel, "has_default_chg_spin") + assert not hasattr(make_base_model(), "has_default_chg_spin") + + def test_absent_from_concrete_models(self) -> None: + for cls in (EnergyModel, SpinModel): + assert not hasattr(cls, "has_default_chg_spin")