feat(pt_expt): full model (#5244)

<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit

* **New Features**
* Added a PyTorch experimental energy model with exportable lower-level
tracing, descriptor/accessor and output-definition APIs, and output-bias
management.

* **Bug Fixes**
* Fixed device placement for created tensors and added runtime
validations for several model accessors.

* **Tests**
* Expanded test suite with autodiff/derivative validation,
export/tracing checks, and cross-backend API consistency tests.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Co-authored-by: Han Wang <wang_han@iapcm.ac.cn>

Author: wanghan-iapcm

deepmd/dpmodel/fitting/general_fitting.py

@@ -584,6 +584,7 @@ def _call_common(
                 )
 
         # calculate the prediction
+        results: dict[str, Array] = {}
         if not self.mixed_types:
             outs = xp.zeros(
                 [nf, nloc, net_dim_out],
@@ -622,4 +623,5 @@ def _call_common(
         exclude_mask = xp.astype(exclude_mask, xp.bool)
         # nf x nloc x nod
         outs = xp.where(exclude_mask[:, :, None], outs, xp.zeros_like(outs))
-        return {self.var_name: outs}
+        results[self.var_name] = outs
+        return results

deepmd/dpmodel/model/dp_model.py

@@ -48,3 +48,7 @@ def update_sel(
     def get_fitting_net(self) -> BaseFitting:
         """Get the fitting network."""
         return self.atomic_model.fitting
+
+    def get_descriptor(self) -> BaseDescriptor:
+        """Get the descriptor."""
+        return self.atomic_model.descriptor

deepmd/dpmodel/model/ener_model.py

@@ -47,3 +47,28 @@ def atomic_output_def(self) -> FittingOutputDef:
         if self._enable_hessian:
             return self.hess_fitting_def
         return super().atomic_output_def()
+
+    def translated_output_def(self) -> dict[str, Any]:
+        """Get the translated output definition.
+
+        Maps internal output names to user-facing names, e.g.
+        ``energy_redu`` -> ``energy``, ``energy_derv_r`` -> ``force``.
+        """
+        out_def_data = self.model_output_def().get_data()
+        output_def = {
+            "atom_energy": out_def_data["energy"],
+            "energy": out_def_data["energy_redu"],
+        }
+        if self.do_grad_r("energy"):
+            output_def["force"] = out_def_data["energy_derv_r"]
+            output_def["force"].squeeze(-2)
+        if self.do_grad_c("energy"):
+            output_def["virial"] = out_def_data["energy_derv_c_redu"]
+            output_def["virial"].squeeze(-2)
+            output_def["atom_virial"] = out_def_data["energy_derv_c"]
+            output_def["atom_virial"].squeeze(-2)
+        if "mask" in out_def_data:
+            output_def["mask"] = out_def_data["mask"]
+        if self._enable_hessian:
+            output_def["hessian"] = out_def_data["energy_derv_r_derv_r"]
+        return output_def

deepmd/dpmodel/model/make_model.py

@@ -21,6 +21,7 @@
     PRECISION_DICT,
     RESERVED_PRECISION_DICT,
     NativeOP,
+    get_xp_precision,
 )
 from deepmd.dpmodel.model.base_model import (
     BaseModel,
@@ -103,7 +104,8 @@ def model_call_from_call_lower(
             bb.reshape(nframes, 3, 3),
         )
     else:
-        coord_normalized = cc.copy()
+        xp = array_api_compat.array_namespace(cc)
+        coord_normalized = xp.reshape(cc, (nframes, nloc, 3))
     extended_coord, extended_atype, mapping = extend_coord_with_ghosts(
         coord_normalized, atype, bb, rcut
     )
@@ -255,7 +257,7 @@ def call(
                 The keys are defined by the `ModelOutputDef`.
 
             """
-            cc, bb, fp, ap, input_prec = self.input_type_cast(
+            cc, bb, fp, ap, input_prec = self._input_type_cast(
                 coord, box=box, fparam=fparam, aparam=aparam
             )
             del coord, box, fparam, aparam
@@ -272,7 +274,7 @@ def call(
                 aparam=ap,
                 do_atomic_virial=do_atomic_virial,
             )
-            model_predict = self.output_type_cast(model_predict, input_prec)
+            model_predict = self._output_type_cast(model_predict, input_prec)
             return model_predict
 
         def call_lower(
@@ -321,7 +323,7 @@ def call_lower(
                 nlist,
                 extra_nlist_sort=self.need_sorted_nlist_for_lower(),
             )
-            cc_ext, _, fp, ap, input_prec = self.input_type_cast(
+            cc_ext, _, fp, ap, input_prec = self._input_type_cast(
                 extended_coord, fparam=fparam, aparam=aparam
             )
             del extended_coord, fparam, aparam
@@ -334,7 +336,7 @@ def call_lower(
                 aparam=ap,
                 do_atomic_virial=do_atomic_virial,
             )
-            model_predict = self.output_type_cast(model_predict, input_prec)
+            model_predict = self._output_type_cast(model_predict, input_prec)
             return model_predict
 
         def forward_common_atomic(
@@ -364,60 +366,107 @@ def forward_common_atomic(
             )
 
         forward_lower = call_lower
+        forward_common = call
+        forward_common_lower = call_lower
 
-        def input_type_cast(
+        def get_out_bias(self) -> Array:
+            """Get the output bias."""
+            return self.atomic_model.out_bias
+
+        def set_out_bias(self, out_bias: Array) -> None:
+            """Set the output bias."""
+            self.atomic_model.out_bias = out_bias
+
+        def change_out_bias(
+            self,
+            merged: Any,
+            bias_adjust_mode: str = "change-by-statistic",
+        ) -> None:
+            """Change the output bias according to the input data and the pretrained model.
+
+            Parameters
+            ----------
+            merged
+                The merged data samples.
+            bias_adjust_mode : str
+                The mode for changing output bias:
+                'change-by-statistic' or 'set-by-statistic'.
+            """
+            self.atomic_model.change_out_bias(merged, bias_adjust_mode=bias_adjust_mode)
+
+        def _input_type_cast(
             self,
             coord: Array,
             box: Array | None = None,
             fparam: Array | None = None,
             aparam: Array | None = None,
-        ) -> tuple[Array, Array, np.ndarray | None, np.ndarray | None, str]:
+        ) -> tuple[Array, Array | None, Array | None, Array | None, Any]:
             """Cast the input data to global float type."""
-            input_prec = RESERVED_PRECISION_DICT[self.precision_dict[coord.dtype.name]]
+            xp = array_api_compat.array_namespace(coord)
+            input_dtype = coord.dtype
+            global_dtype = get_xp_precision(
+                xp, RESERVED_PRECISION_DICT[self.global_np_float_precision]
+            )
             ###
             ### type checking would not pass jit, convert to coord prec anyway
             ###
-            _lst: list[np.ndarray | None] = [
-                vv.astype(coord.dtype) if vv is not None else None
+            _lst: list[Array | None] = [
+                xp.astype(vv, input_dtype) if vv is not None else None
                 for vv in [box, fparam, aparam]
             ]
             box, fparam, aparam = _lst
-            if input_prec == RESERVED_PRECISION_DICT[self.global_np_float_precision]:
-                return coord, box, fparam, aparam, input_prec
+            if input_dtype == global_dtype:
+                return coord, box, fparam, aparam, input_dtype
             else:
-                pp = self.global_np_float_precision
                 return (
-                    coord.astype(pp),
-                    box.astype(pp) if box is not None else None,
-                    fparam.astype(pp) if fparam is not None else None,
-                    aparam.astype(pp) if aparam is not None else None,
-                    input_prec,
+                    xp.astype(coord, global_dtype),
+                    xp.astype(box, global_dtype) if box is not None else None,
+                    xp.astype(fparam, global_dtype) if fparam is not None else None,
+                    xp.astype(aparam, global_dtype) if aparam is not None else None,
+                    input_dtype,
                 )
 
-        def output_type_cast(
+        def _output_type_cast(
             self,
             model_ret: dict[str, Array],
-            input_prec: str,
+            input_prec: Any,
         ) -> dict[str, Array]:
-            """Convert the model output to the input prec."""
-            do_cast = (
-                input_prec != RESERVED_PRECISION_DICT[self.global_np_float_precision]
+            """Convert the model output to the input prec.
+
+            Parameters
+            ----------
+            model_ret
+                The model output.
+            input_prec
+                The input dtype returned by ``_input_type_cast``.
+            """
+            model_ret_not_none = [vv for vv in model_ret.values() if vv is not None]
+            if not model_ret_not_none:
+                return model_ret
+            xp = array_api_compat.array_namespace(model_ret_not_none[0])
+            global_dtype = get_xp_precision(
+                xp, RESERVED_PRECISION_DICT[self.global_np_float_precision]
+            )
+            ener_dtype = get_xp_precision(
+                xp, RESERVED_PRECISION_DICT[self.global_ener_float_precision]
             )
-            pp = self.precision_dict[input_prec]
+            do_cast = input_prec != global_dtype
             odef = self.model_output_def()
             for kk in odef.keys():
                 if kk not in model_ret.keys():
                     # do not return energy_derv_c if not do_atomic_virial
                     continue
                 if check_operation_applied(odef[kk], OutputVariableOperation.REDU):
                     model_ret[kk] = (
-                        model_ret[kk].astype(self.global_ener_float_precision)
+                        xp.astype(model_ret[kk], ener_dtype)
                         if model_ret[kk] is not None
                         else None
                     )
                 elif do_cast:
                     model_ret[kk] = (
-                        model_ret[kk].astype(pp) if model_ret[kk] is not None else None
+                        xp.astype(model_ret[kk], input_prec)
+                        if model_ret[kk] is not None
+                        else None
                     )
             return model_ret
 

deepmd/dpmodel/model/transform_output.py

@@ -98,6 +98,7 @@ def communicate_extended_output(
 
     """
     xp = array_api_compat.get_namespace(mapping)
+    device = array_api_compat.device(mapping)
     mapping_ = mapping
     new_ret = {}
     for kk in model_output_def.keys_outp():
@@ -117,7 +118,9 @@ def communicate_extended_output(
                         mapping, tuple(mldims + [1] * len(derv_r_ext_dims))
                     )
                     mapping = xp.tile(mapping, [1] * len(mldims) + derv_r_ext_dims)
-                    force = xp.zeros(vldims + derv_r_ext_dims, dtype=vv.dtype)
+                    force = xp.zeros(
+                        vldims + derv_r_ext_dims, dtype=vv.dtype, device=device
+                    )
                     force = xp_scatter_sum(
                         force,
                         1,
@@ -149,7 +152,9 @@ def communicate_extended_output(
                         nall = hess_1.shape[1]
                         # (1) -> [nf, nloc1, nall2, *def, 3(1), 3(2)]
                         hessian1 = xp.zeros(
-                            [*vldims, nall, *vdef.shape, 3, 3], dtype=vv.dtype
+                            [*vldims, nall, *vdef.shape, 3, 3],
+                            dtype=vv.dtype,
+                            device=device,
                         )
                         mapping_hess = xp.reshape(
                             mapping_, (mldims + [1] * (len(vdef.shape) + 3))
@@ -172,7 +177,9 @@ def communicate_extended_output(
                         nloc = hessian1.shape[2]
                         # (2) -> [nf, nloc2, nloc1, *def, 3(1), 3(2)]
                         hessian = xp.zeros(
-                            [*vldims, nloc, *vdef.shape, 3, 3], dtype=vv.dtype
+                            [*vldims, nloc, *vdef.shape, 3, 3],
+                            dtype=vv.dtype,
+                            device=device,
                         )
                         mapping_hess = xp.reshape(
                             mapping_, (mldims + [1] * (len(vdef.shape) + 3))
@@ -218,6 +225,7 @@ def communicate_extended_output(
                     virial = xp.zeros(
                         vldims + derv_c_ext_dims,
                         dtype=vv.dtype,
+                        device=device,
                     )
                     virial = xp_scatter_sum(
                         virial,

deepmd/dpmodel/utils/network.py

@@ -280,11 +280,11 @@ def call(self, x):  # noqa: ANN001, ANN201
             y = xp.astype(y, x.dtype)
         y = fn(y)
         if self.idt is not None:
-            y *= self.idt
+            y = y * self.idt
         if self.resnet and self.w.shape[1] == self.w.shape[0]:
-            y += x
+            y = y + x
         elif self.resnet and self.w.shape[1] == 2 * self.w.shape[0]:
-            y += xp.concat([x, x], axis=-1)
+            y = y + xp.concat([x, x], axis=-1)
         return y
 
 

deepmd/pd/model/model/ener_model.py

@@ -60,7 +60,7 @@ def translated_output_def(self) -> dict:
             output_def["virial"] = out_def_data["energy_derv_c_redu"]
             output_def["virial"].squeeze(-2)
             output_def["atom_virial"] = out_def_data["energy_derv_c"]
-            output_def["atom_virial"].squeeze(-3)
+            output_def["atom_virial"].squeeze(-2)
         if "mask" in out_def_data:
             output_def["mask"] = out_def_data["mask"]
         return output_def
@@ -140,7 +140,7 @@ def forward_lower(
                 if do_atomic_virial:
                     model_predict["extended_virial"] = model_ret[
                         "energy_derv_c"
-                    ].squeeze(-3)
+                    ].squeeze(-2)
                 else:
                     model_predict["extended_virial"] = paddle.zeros(
                         [model_predict["energy"].shape[0], 1, 9], dtype=paddle.float64

deepmd/pd/model/model/make_model.py

@@ -162,7 +162,7 @@ def forward_common(
                 The keys are defined by the `ModelOutputDef`.
 
             """
-            cc, bb, fp, ap, input_prec = self.input_type_cast(
+            cc, bb, fp, ap, input_prec = self._input_type_cast(
                 coord, box=box, fparam=fparam, aparam=aparam
             )
             del coord, box, fparam, aparam
@@ -196,7 +196,7 @@ def forward_common(
                 mapping,
                 do_atomic_virial=do_atomic_virial,
             )
-            model_predict = self.output_type_cast(model_predict, input_prec)
+            model_predict = self._output_type_cast(model_predict, input_prec)
             return model_predict
 
         def get_out_bias(self) -> paddle.Tensor:
@@ -283,7 +283,7 @@ def forward_common_lower(
             nlist = self.format_nlist(
                 extended_coord, extended_atype, nlist, extra_nlist_sort=extra_nlist_sort
             )
-            cc_ext, _, fp, ap, input_prec = self.input_type_cast(
+            cc_ext, _, fp, ap, input_prec = self._input_type_cast(
                 extended_coord, fparam=fparam, aparam=aparam
             )
             del extended_coord, fparam, aparam
@@ -303,10 +303,10 @@ def forward_common_lower(
                 do_atomic_virial=do_atomic_virial,
                 create_graph=self.training,
             )
-            model_predict = self.output_type_cast(model_predict, input_prec)
+            model_predict = self._output_type_cast(model_predict, input_prec)
             return model_predict
 
-        def input_type_cast(
+        def _input_type_cast(
             self,
             coord: paddle.Tensor,
             box: paddle.Tensor | None = None,
@@ -351,7 +351,7 @@ def input_type_cast(
                     input_prec,
                 )
 
-        def output_type_cast(
+        def _output_type_cast(
             self,
             model_ret: dict[str, paddle.Tensor],
             input_prec: str,

deepmd/pt/model/atomic_model/dp_atomic_model.py

@@ -83,6 +83,11 @@ def set_eval_descriptor_hook(self, enable: bool) -> None:
 
     def eval_descriptor(self) -> torch.Tensor:
         """Evaluate the descriptor."""
+        if not self.eval_descriptor_list:
+            raise RuntimeError(
+                "eval_descriptor_list is empty. "
+                "Call set_eval_descriptor_hook(True) and perform a forward pass first."
+            )
         return torch.concat(self.eval_descriptor_list)
 
     def set_eval_fitting_last_layer_hook(self, enable: bool) -> None:
@@ -94,6 +99,11 @@ def set_eval_fitting_last_layer_hook(self, enable: bool) -> None:
 
     def eval_fitting_last_layer(self) -> torch.Tensor:
         """Evaluate the fitting last layer output."""
+        if not self.eval_fitting_last_layer_list:
+            raise RuntimeError(
+                "eval_fitting_last_layer_list is empty. "
+                "Call set_eval_fitting_last_layer_hook(True) and perform a forward pass first."
+            )
         return torch.concat(self.eval_fitting_last_layer_list)
 
     @torch.jit.export