fix: cast input and output types at model's interface

deepmd/pt/model/model/make_model.py

@@ -3,6 +3,7 @@
     Dict,
     List,
     Optional,
+    Tuple,
 )
 
 import torch
@@ -17,6 +18,11 @@
     communicate_extended_output,
     fit_output_to_model_output,
 )
+from deepmd.pt.utils.env import (
+    GLOBAL_PT_FLOAT_PRECISION,
+    PRECISION_DICT,
+    RESERVED_PRECISON_DICT,
+)
 from deepmd.pt.utils.nlist import (
     extend_input_and_build_neighbor_list,
     nlist_distinguish_types,
@@ -56,6 +62,9 @@
                 *args,
                 **kwargs,
             )
+            self.precision_dict = PRECISION_DICT
+            self.reverse_precision_dict = RESERVED_PRECISON_DICT
+            self.global_pt_float_precision = GLOBAL_PT_FLOAT_PRECISION
 
         def model_output_def(self):
             """Get the output def for the model."""
@@ -115,34 +124,39 @@
                 The keys are defined by the `ModelOutputDef`.
 
             """
+            cc, bb, fp, ap, input_prec = self.input_type_cast(
+                coord, box=box, fparam=fparam, aparam=aparam
+            )
+            del coord, box, fparam, aparam
             (
                 extended_coord,
                 extended_atype,
                 mapping,
                 nlist,
             ) = extend_input_and_build_neighbor_list(
-                coord,
+                cc,
                 atype,
                 self.get_rcut(),
                 self.get_sel(),
                 mixed_types=self.mixed_types(),
-                box=box,
+                box=bb,
             )
             model_predict_lower = self.forward_common_lower(
                 extended_coord,
                 extended_atype,
                 nlist,
                 mapping,
                 do_atomic_virial=do_atomic_virial,
-                fparam=fparam,
-                aparam=aparam,
+                fparam=fp,
+                aparam=ap,
             )
             model_predict = communicate_extended_output(
                 model_predict_lower,
                 self.model_output_def(),
                 mapping,
                 do_atomic_virial=do_atomic_virial,
             )
+            model_predict = self.output_type_cast(model_predict, input_prec)
             return model_predict
 
         def forward_common_lower(
@@ -186,22 +200,87 @@
             nframes, nall = extended_atype.shape[:2]
             extended_coord = extended_coord.view(nframes, -1, 3)
             nlist = self.format_nlist(extended_coord, extended_atype, nlist)
+            cc_ext, _, fp, ap, input_prec = self.input_type_cast(
+                extended_coord, fparam=fparam, aparam=aparam
+            )
+            del extended_coord, fparam, aparam
             atomic_ret = self.forward_atomic(
-                extended_coord,
+                cc_ext,
                 extended_atype,
                 nlist,
                 mapping=mapping,
-                fparam=fparam,
-                aparam=aparam,
+                fparam=fp,
+                aparam=ap,
             )
             model_predict = fit_output_to_model_output(
                 atomic_ret,
                 self.fitting_output_def(),
-                extended_coord,
+                cc_ext,
                 do_atomic_virial=do_atomic_virial,
             )
+            model_predict = self.output_type_cast(model_predict, input_prec)
             return model_predict
 
+        def input_type_cast(
+            self,
+            coord: torch.Tensor,
+            box: Optional[torch.Tensor] = None,
+            fparam: Optional[torch.Tensor] = None,
+            aparam: Optional[torch.Tensor] = None,
+        ) -> Tuple[
+            torch.Tensor,
+            Optional[torch.Tensor],
+            Optional[torch.Tensor],
+            Optional[torch.Tensor],
+            str,
+        ]:
+            """Cast the input data to global float type."""
+            input_prec = self.reverse_precision_dict[coord.dtype]
+            ###
+            ### type checking would not pass jit, convert to coord prec anyway
+            ###
+            # for vv, kk in zip([fparam, aparam], ["frame", "atomic"]):
+            #     if vv is not None and self.reverse_precision_dict[vv.dtype] != input_prec:
+            #         log.warning(
+            #           f"type of {kk} parameter {self.reverse_precision_dict[vv.dtype]}"
+            #           " does not match"
+            #           f" that of the coordinate {input_prec}"
+            #         )
+            _lst: List[Optional[torch.Tensor]] = [
+                vv.to(coord.dtype) if vv is not None else None
+                for vv in [box, fparam, aparam]
+            ]
+            box, fparam, aparam = _lst
+            if (
+                input_prec
+                == self.reverse_precision_dict[self.global_pt_float_precision]
+            ):
+                return coord, box, fparam, aparam, input_prec
+            else:
+                pp = self.global_pt_float_precision
+                return (
+                    coord.to(pp),
+                    box.to(pp) if box is not None else None,
+                    fparam.to(pp) if fparam is not None else None,
+                    aparam.to(pp) if aparam is not None else None,
+                    input_prec,
+                )
+
+        def output_type_cast(
+            self,
+            model_ret: Dict[str, torch.Tensor],
+            input_prec: str,
+        ) -> Dict[str, torch.Tensor]:
+            """Convert the model output to the input prec."""
+            if (
+                input_prec
+                != self.reverse_precision_dict[self.global_pt_float_precision]
+            ):
+                pp = self.precision_dict[input_prec]
+                for kk, vv in model_ret.items():
+                    model_ret[kk] = vv.to(pp) if vv is not None else None
+            return model_ret
+
         @torch.jit.export
         def format_nlist(
             self,

source/tests/pt/model/test_dp_model.py

@@ -186,6 +186,49 @@ def test_dp_consistency_nopbc(self):
             to_numpy_array(ret1["energy_redu"]),
         )
 
+    def test_prec_consistency(self):
+        rng = np.random.default_rng()
+        nf, nloc = self.atype.shape
+        ds = DPDescrptSeA(
+            self.rcut,
+            self.rcut_smth,
+            self.sel,
+        )
+        ft = DPInvarFitting(
+            "energy",
+            self.nt,
+            ds.get_dim_out(),
+            1,
+            mixed_types=ds.mixed_types(),
+        )
+        nfp, nap = 2, 3
+        type_map = ["foo", "bar"]
+        fparam = rng.normal(size=[self.nf, nfp])
+        aparam = rng.normal(size=[self.nf, nloc, nap])
+
+        md0 = DPDPModel(ds, ft, type_map=type_map)
+        md1 = DPModel.deserialize(md0.serialize()).to(env.DEVICE)
+
+        args64 = [to_torch_tensor(ii) for ii in [self.coord, self.atype, self.cell]]
+        args64[0] = args64[0].to(torch.float64)
+        args64[2] = args64[2].to(torch.float64)
+        args32 = [to_torch_tensor(ii) for ii in [self.coord, self.atype, self.cell]]
+        args32[0] = args32[0].to(torch.float32)
+        args32[2] = args32[2].to(torch.float32)
+        # fparam, aparam are converted to coordinate precision by model
+        fparam = to_torch_tensor(fparam)
+        aparam = to_torch_tensor(aparam)
+
+        model_l_ret_64 = md1.forward_common(*args64, fparam=fparam, aparam=aparam)
+        model_l_ret_32 = md1.forward_common(*args32, fparam=fparam, aparam=aparam)
+
+        for ii in model_l_ret_32.keys():
+            self.assertEqual(model_l_ret_32[ii].dtype, torch.float32)
+            np.testing.assert_allclose(
+                to_numpy_array(model_l_ret_32[ii]),
+                to_numpy_array(model_l_ret_64[ii]),
+            )
+
 
 class TestDPModelLower(unittest.TestCase, TestCaseSingleFrameWithNlist):
     def setUp(self):
@@ -269,6 +312,51 @@ def test_dp_consistency(self):
             to_numpy_array(ret1["energy_redu"]),
         )
 
+    def test_prec_consistency(self):
+        rng = np.random.default_rng()
+        nf, nloc, nnei = self.nlist.shape
+        ds = DPDescrptSeA(
+            self.rcut,
+            self.rcut_smth,
+            self.sel,
+        )
+        ft = DPInvarFitting(
+            "energy",
+            self.nt,
+            ds.get_dim_out(),
+            1,
+            mixed_types=ds.mixed_types(),
+        )
+        nfp, nap = 2, 3
+        type_map = ["foo", "bar"]
+        fparam = rng.normal(size=[self.nf, nfp])
+        aparam = rng.normal(size=[self.nf, nloc, nap])
+
+        md0 = DPDPModel(ds, ft, type_map=type_map)
+        md1 = DPModel.deserialize(md0.serialize()).to(env.DEVICE)
+
+        args64 = [
+            to_torch_tensor(ii) for ii in [self.coord_ext, self.atype_ext, self.nlist]
+        ]
+        args64[0] = args64[0].to(torch.float64)
+        args32 = [
+            to_torch_tensor(ii) for ii in [self.coord_ext, self.atype_ext, self.nlist]
+        ]
+        args32[0] = args32[0].to(torch.float32)
+        # fparam, aparam are converted to coordinate precision by model
+        fparam = to_torch_tensor(fparam)
+        aparam = to_torch_tensor(aparam)
+
+        model_l_ret_64 = md1.forward_common_lower(*args64, fparam=fparam, aparam=aparam)
+        model_l_ret_32 = md1.forward_common_lower(*args32, fparam=fparam, aparam=aparam)
+
+        for ii in model_l_ret_32.keys():
+            self.assertEqual(model_l_ret_32[ii].dtype, torch.float32)
+            np.testing.assert_allclose(
+                to_numpy_array(model_l_ret_32[ii]),
+                to_numpy_array(model_l_ret_64[ii]),
+            )
+
     def test_jit(self):
         nf, nloc, nnei = self.nlist.shape
         ds = DescrptSeA(