fix(pt_expt): dispatch property models to DeepProperty inference

deepmd/dpmodel/infer/deep_eval.py

@@ -47,6 +47,9 @@
 from deepmd.infer.deep_pot import (
     DeepPot,
 )
+from deepmd.infer.deep_property import (
+    DeepProperty,
+)
 from deepmd.infer.deep_wfc import (
     DeepWFC,
 )
@@ -134,8 +137,9 @@ def has_default_fparam(self) -> bool:
 
     @property
     def model_type(self) -> type["DeepEvalWrapper"]:
-        """The the evaluator of the model type."""
-        model_output_type = self.dp.model_output_type()
+        """The evaluator of the model type."""
+        model = self.get_model()
+        model_output_type = model.model_output_type()
         if "energy" in model_output_type:
             return DeepPot
         elif "dos" in model_output_type:
@@ -146,6 +150,8 @@ def model_type(self) -> type["DeepEvalWrapper"]:
             return DeepPolar
         elif "wfc" in model_output_type:
             return DeepWFC
+        elif self._get_property_var_name(model) in model_output_type:
+            return DeepProperty
         else:
             raise RuntimeError("Unknown model type")
 
@@ -238,6 +244,12 @@ def eval(
         out = self._eval_func(self._eval_model, numb_test, natoms)(
             coords, cells, atom_types, fparam, aparam, request_defs
         )
+        # ``AutoBatchSize.execute_all`` unwraps a single-output result out of
+        # its tuple, which would make ``zip`` iterate over the array's frame
+        # axis. Re-wrap so the request-def names line up (a single request def
+        # arises for global-only DOS/property inference at atomic=False).
+        if not isinstance(out, tuple):
+            out = (out,)
         return dict(
             zip(
                 [x.name for x in request_defs],

deepmd/dpmodel/model/property_model.py

@@ -43,6 +43,14 @@ def get_var_name(self) -> str:
         """Get the name of the property."""
         return self.get_fitting_net().var_name
 
+    def get_task_dim(self) -> int:
+        """Get the output dimension of the property."""
+        return self.get_fitting_net().dim_out
+
+    def get_intensive(self) -> bool:
+        """Whether the property is intensive."""
+        return self.model_output_def()[self.get_var_name()].intensive
+
     def call(
         self,
         coord: Array,

deepmd/infer/deep_eval.py

@@ -375,7 +375,31 @@ def _check_mixed_types(self, atom_types: np.ndarray) -> bool:
     @property
     @abstractmethod
     def model_type(self) -> type["DeepEval"]:
-        """The the evaluator of the model type."""
+        """The evaluator of the model type.
+
+        Each backend implements the dispatch on its own module so it can import
+        the concrete ``Deep*`` wrapper classes at the top level. Those wrappers
+        import ``DeepEval`` from this module, so a dispatch here would form an
+        import cycle (flagged by CodeQL). :meth:`_get_property_var_name` is
+        provided for the shared property branch.
+        """
+
+    @staticmethod
+    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.
+        """
+        if not hasattr(model, "get_var_name"):
+            return None
+        try:
+            return model.get_var_name()
+        except NotImplementedError:
+            return None
 
     @abstractmethod
     def get_sel_type(self) -> list[int]:
@@ -414,7 +438,24 @@ def get_has_hessian(self) -> bool:
         return False
 
     def get_var_name(self) -> str:
-        """Get the name of the fitting property."""
+        """Get the name of the fitting property (property models only)."""
+        model = self.get_model()
+        if hasattr(model, "get_var_name"):
+            return model.get_var_name()
+        raise NotImplementedError
+
+    def get_task_dim(self) -> int:
+        """Get the output dimension of the property (property models only)."""
+        model = self.get_model()
+        if hasattr(model, "get_task_dim"):
+            return model.get_task_dim()
+        raise NotImplementedError
+
+    def get_intensive(self) -> bool:
+        """Whether the property is intensive (property models only)."""
+        model = self.get_model()
+        if hasattr(model, "get_intensive"):
+            return model.get_intensive()
         raise NotImplementedError
 
     @abstractmethod

deepmd/infer/deep_property.py

@@ -134,14 +134,13 @@ def eval(
             aparam=aparam,
             **kwargs,
         )
-        atomic_property = results[self.get_var_name()].reshape(
-            nframes, natoms, self.get_task_dim()
-        )
         property = results[f"{self.get_var_name()}_redu"].reshape(
             nframes, self.get_task_dim()
         )
-
         if atomic:
+            atomic_property = results[self.get_var_name()].reshape(
+                nframes, natoms, self.get_task_dim()
+            )
             return (
                 property,
                 atomic_property,

deepmd/jax/infer/deep_eval.py

@@ -41,6 +41,9 @@
 from deepmd.infer.deep_pot import (
     DeepPot,
 )
+from deepmd.infer.deep_property import (
+    DeepProperty,
+)
 from deepmd.infer.deep_wfc import (
     DeepWFC,
 )
@@ -152,7 +155,8 @@ def get_dim_aparam(self) -> int:
     @property
     def model_type(self) -> type["DeepEvalWrapper"]:
         """The evaluator of the model type."""
-        model_output_type = self.dp.model_output_type()
+        model = self.get_model()
+        model_output_type = model.model_output_type()
         if "energy" in model_output_type:
             return DeepPot
         elif "dos" in model_output_type:
@@ -163,6 +167,8 @@ def model_type(self) -> type["DeepEvalWrapper"]:
             return DeepPolar
         elif "wfc" in model_output_type:
             return DeepWFC
+        elif self._get_property_var_name(model) in model_output_type:
+            return DeepProperty
         else:
             raise RuntimeError("Unknown model type")
 
@@ -270,6 +276,12 @@ def eval(
         out = self._eval_func(self._eval_model, numb_test, natoms)(
             coords, cells, atom_types, fparam, aparam, request_defs
         )
+        # ``AutoBatchSize.execute_all`` unwraps a single-output result out of
+        # its tuple, which would make ``zip`` iterate over the array's frame
+        # axis. Re-wrap so the request-def names line up (a single request def
+        # arises for global-only DOS/property inference at atomic=False).
+        if not isinstance(out, tuple):
+            out = (out,)
         return dict(
             zip(
                 [x.name for x in request_defs],

deepmd/jax/jax2tf/serialization.py

@@ -318,6 +318,28 @@ def get_default_fparam() -> tf.Tensor:
 
         tf_model.get_default_fparam = get_default_fparam
 
+        # 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"):
+
+            @tf.function
+            def get_var_name() -> tf.Tensor:
+                return tf.constant(model.get_var_name(), dtype=tf.string)
+
+            tf_model.get_var_name = get_var_name
+
+            @tf.function
+            def get_task_dim() -> tf.Tensor:
+                return tf.constant(model.get_task_dim(), dtype=tf.int64)
+
+            tf_model.get_task_dim = get_task_dim
+
+            @tf.function
+            def get_intensive() -> tf.Tensor:
+                return tf.constant(model.get_intensive(), dtype=tf.bool)
+
+            tf_model.get_intensive = get_intensive
+
         tf.saved_model.save(
             tf_model,
             model_file,

deepmd/jax/jax2tf/tfmodel.py

@@ -78,6 +78,15 @@ def __init__(
             self.default_fparam = self.model.get_default_fparam().numpy().tolist()
         else:
             self.default_fparam = None
+        # property models only (absent for other model types).
+        if hasattr(self.model, "get_var_name"):
+            self._var_name = self.model.get_var_name().numpy().decode()
+            self._task_dim = self.model.get_task_dim().numpy().item()
+            self._intensive = self.model.get_intensive().numpy().item()
+        else:
+            self._var_name = None
+            self._task_dim = None
+            self._intensive = False
 
     def __call__(
         self,
@@ -175,9 +184,27 @@ def call(
 
     def model_output_def(self) -> ModelOutputDef:
         return ModelOutputDef(
-            FittingOutputDef([OUTPUT_DEFS[tt] for tt in self.model_output_type()])
+            FittingOutputDef(
+                [self._output_var_def(tt) for tt in self.model_output_type()]
+            )
         )
 
+    def _output_var_def(self, name: str) -> OutputVariableDef:
+        if name in OUTPUT_DEFS:
+            return OUTPUT_DEFS[name]
+        # property models carry a user-defined output name (``var_name``) that
+        # is not in the fixed table; rebuild its def from the persisted metadata.
+        if self._var_name is not None and name == self._var_name:
+            return OutputVariableDef(
+                self._var_name,
+                shape=[self._task_dim],
+                reducible=True,
+                r_differentiable=False,
+                c_differentiable=False,
+                intensive=self._intensive,
+            )
+        raise KeyError(f"Unknown model output variable {name!r}")
+
     def call_lower(
         self,
         extended_coord: jnp.ndarray,
@@ -349,3 +376,19 @@ def has_default_fparam(self) -> bool:
     def get_default_fparam(self) -> list[float] | None:
         """Get the default frame parameters."""
         return self.default_fparam
+
+    def get_var_name(self) -> str:
+        """Get the name of the property (property models only)."""
+        if self._var_name is None:
+            raise NotImplementedError
+        return self._var_name
+
+    def get_task_dim(self) -> int:
+        """Get the output dimension of the property (property models only)."""
+        if self._task_dim is None:
+            raise NotImplementedError
+        return self._task_dim
+
+    def get_intensive(self) -> bool:
+        """Whether the property is intensive (property models only)."""
+        return self._intensive

deepmd/jax/model/hlo.py

@@ -69,6 +69,10 @@ def __init__(
         has_default_fparam: bool = False,
         default_fparam: list[float] | None = None,
         numb_dos: int = 0,
+        # property models only
+        var_name: str | None = None,
+        task_dim: int | None = None,
+        intensive: bool = False,
     ) -> None:
         self._call_lower = jax_export.deserialize(stablehlo).call
         self._call_lower_atomic_virial = jax_export.deserialize(
@@ -93,6 +97,9 @@ def __init__(
         self._has_default_fparam = has_default_fparam
         self.default_fparam = default_fparam
         self.numb_dos = numb_dos
+        self._var_name = var_name
+        self._task_dim = task_dim
+        self._intensive = intensive
 
     def __call__(
         self,
@@ -180,9 +187,27 @@ def call(
 
     def model_output_def(self) -> ModelOutputDef:
         return ModelOutputDef(
-            FittingOutputDef([OUTPUT_DEFS[tt] for tt in self.model_output_type()])
+            FittingOutputDef(
+                [self._output_var_def(tt) for tt in self.model_output_type()]
+            )
         )
 
+    def _output_var_def(self, name: str) -> OutputVariableDef:
+        if name in OUTPUT_DEFS:
+            return OUTPUT_DEFS[name]
+        # property models carry a user-defined output name (``var_name``) that
+        # is not in the fixed table; rebuild its def from the persisted metadata.
+        if self._var_name is not None and name == self._var_name:
+            return OutputVariableDef(
+                self._var_name,
+                shape=[self._task_dim],
+                reducible=True,
+                r_differentiable=False,
+                c_differentiable=False,
+                intensive=self._intensive,
+            )
+        raise KeyError(f"Unknown model output variable {name!r}")
+
     def call_lower(
         self,
         extended_coord: jnp.ndarray,
@@ -233,6 +258,22 @@ def get_dim_aparam(self) -> int:
         """Get the number (dimension) of atomic parameters of this atomic model."""
         return self.dim_aparam
 
+    def get_var_name(self) -> str:
+        """Get the name of the property (property models only)."""
+        if self._var_name is None:
+            raise NotImplementedError
+        return self._var_name
+
+    def get_task_dim(self) -> int:
+        """Get the output dimension of the property (property models only)."""
+        if self._task_dim is None:
+            raise NotImplementedError
+        return self._task_dim
+
+    def get_intensive(self) -> bool:
+        """Whether the property is intensive (property models only)."""
+        return self._intensive
+
     def get_sel_type(self) -> list[int]:
         """Get the selected atom types of this model.
 

deepmd/jax/utils/serialization.py

@@ -321,6 +321,18 @@ def call_lower_with_fixed_do_atomic_virial(
             "sel": model.get_sel(),
             "has_default_fparam": model.has_default_fparam(),
             "default_fparam": model.get_default_fparam(),
+            # 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,
         }
         save_dp_model(filename=model_file, model_dict=data)
     elif model_file.endswith(".savedmodel"):