fix(np/pt): explicit dtype and device. (#4241)

fix: #4016

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

## Release Notes

- **New Features**
- Enhanced precision handling for various models and computations,
ensuring consistent data types across multiple components.
- Introduced new parameters for model fitting processes to improve
flexibility and control over data handling.

- **Bug Fixes**
- Improved error handling in several methods to prevent runtime issues
related to data type mismatches and invalid configurations.

- **Documentation**
- Updated comments and documentation strings for clarity on parameter
usage and expected behaviors.

- **Refactor**
- Streamlined code by enhancing type safety and consistency in tensor
operations, particularly in statistical computations and model fitting
processes.

- **Tests**
- Improved test structure for better maintainability without altering
existing functionality.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Signed-off-by: Han Wang <[email protected]>
Co-authored-by: Han Wang <[email protected]>
Co-authored-by: coderabbitai[bot] <136622811+coderabbitai[bot]@users.noreply.github.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

deepmd/dpmodel/atomic_model/base_atomic_model.py

@@ -19,6 +19,9 @@
     AtomExcludeMask,
     PairExcludeMask,
 )
+from deepmd.env import (
+    GLOBAL_NP_FLOAT_PRECISION,
+)
 from deepmd.utils.finetune import (
     get_index_between_two_maps,
     map_atom_exclude_types,
@@ -56,8 +59,12 @@ def init_out_stat(self):
             [self.atomic_output_def()[kk].size for kk in self.bias_keys]
         )
         self.n_out = len(self.bias_keys)
-        out_bias_data = np.zeros([self.n_out, ntypes, self.max_out_size])  # pylint: disable=no-explicit-dtype
-        out_std_data = np.ones([self.n_out, ntypes, self.max_out_size])  # pylint: disable=no-explicit-dtype
+        out_bias_data = np.zeros(
+            [self.n_out, ntypes, self.max_out_size], dtype=GLOBAL_NP_FLOAT_PRECISION
+        )
+        out_std_data = np.ones(
+            [self.n_out, ntypes, self.max_out_size], dtype=GLOBAL_NP_FLOAT_PRECISION
+        )
         self.out_bias = out_bias_data
         self.out_std = out_std_data
 

deepmd/dpmodel/atomic_model/linear_atomic_model.py

@@ -12,6 +12,9 @@
     get_multiple_nlist_key,
     nlist_distinguish_types,
 )
+from deepmd.env import (
+    GLOBAL_NP_FLOAT_PRECISION,
+)
 from deepmd.utils.version import (
     check_version_compatibility,
 )
@@ -286,7 +289,11 @@ def _compute_weight(
         """This should be a list of user defined weights that matches the number of models to be combined."""
         nmodels = len(self.models)
         nframes, nloc, _ = nlists_[0].shape
-        return [np.ones((nframes, nloc, 1)) / nmodels for _ in range(nmodels)]  # pylint: disable=no-explicit-dtype
+        # the dtype of weights is the interface data type.
+        return [
+            np.ones((nframes, nloc, 1), dtype=GLOBAL_NP_FLOAT_PRECISION) / nmodels
+            for _ in range(nmodels)
+        ]
 
     def get_dim_fparam(self) -> int:
         """Get the number (dimension) of frame parameters of this atomic model."""

deepmd/dpmodel/atomic_model/pairtab_atomic_model.py

@@ -204,9 +204,10 @@ def forward_atomic(
             self.tab.ntypes, self.tab.ntypes, self.tab.nspline, 4
         )
 
-        # (nframes, nloc, nnei)
+        # (nframes, nloc, nnei), index type is int64.
         j_type = extended_atype[
-            np.arange(extended_atype.shape[0])[:, None, None], masked_nlist  # pylint: disable=no-explicit-dtype
+            np.arange(extended_atype.shape[0], dtype=np.int64)[:, None, None],
+            masked_nlist,
         ]
 
         raw_atomic_energy = self._pair_tabulated_inter(
@@ -303,7 +304,8 @@ def _get_pairwise_dist(coords: np.ndarray, nlist: np.ndarray) -> np.ndarray:
         np.ndarray
             The pairwise distance between the atoms (nframes, nloc, nnei).
         """
-        batch_indices = np.arange(nlist.shape[0])[:, None, None]  # pylint: disable=no-explicit-dtype
+        # index type is int64
+        batch_indices = np.arange(nlist.shape[0], dtype=np.int64)[:, None, None]
         neighbor_atoms = coords[batch_indices, nlist]
         loc_atoms = coords[:, : nlist.shape[1], :]
         pairwise_dr = loc_atoms[:, :, None, :] - neighbor_atoms

deepmd/dpmodel/common.py

@@ -50,6 +50,33 @@
 DEFAULT_PRECISION = "float64"
 
 
+def get_xp_precision(
+    xp: Any,
+    precision: str,
+):
+    """Get the precision from the API compatible namespace."""
+    if precision == "float16" or precision == "half":
+        return xp.float16
+    elif precision == "float32" or precision == "single":
+        return xp.float32
+    elif precision == "float64" or precision == "double":
+        return xp.float64
+    elif precision == "int32":
+        return xp.int32
+    elif precision == "int64":
+        return xp.int64
+    elif precision == "bool":
+        return bool
+    elif precision == "default":
+        return get_xp_precision(xp, RESERVED_PRECISON_DICT[PRECISION_DICT[precision]])
+    elif precision == "global":
+        return get_xp_precision(xp, RESERVED_PRECISON_DICT[GLOBAL_NP_FLOAT_PRECISION])
+    elif precision == "bfloat16":
+        return ml_dtypes.bfloat16
+    else:
+        raise ValueError(f"unsupported precision {precision} for {xp}")
+
+
 class NativeOP(ABC):
     """The unit operation of a native model."""
 

deepmd/dpmodel/fitting/general_fitting.py

@@ -14,9 +14,11 @@
 
 from deepmd.dpmodel import (
     DEFAULT_PRECISION,
+    PRECISION_DICT,
     NativeOP,
 )
 from deepmd.dpmodel.common import (
+    get_xp_precision,
     to_numpy_array,
 )
 from deepmd.dpmodel.utils import (
@@ -27,6 +29,9 @@
 from deepmd.dpmodel.utils.seed import (
     child_seed,
 )
+from deepmd.env import (
+    GLOBAL_NP_FLOAT_PRECISION,
+)
 from deepmd.utils.finetune import (
     get_index_between_two_maps,
     map_atom_exclude_types,
@@ -133,6 +138,11 @@ def __init__(
             self.trainable = [self.trainable] * (len(self.neuron) + 1)
         self.activation_function = activation_function
         self.precision = precision
+        if self.precision.lower() not in PRECISION_DICT:
+            raise ValueError(
+                f"Unsupported precision '{self.precision}'. Supported options are: {list(PRECISION_DICT.keys())}"
+            )
+        self.prec = PRECISION_DICT[self.precision.lower()]
         self.layer_name = layer_name
         self.use_aparam_as_mask = use_aparam_as_mask
         self.spin = spin
@@ -146,18 +156,20 @@ def __init__(
         net_dim_out = self._net_out_dim()
         # init constants
         if bias_atom_e is None:
-            self.bias_atom_e = np.zeros([self.ntypes, net_dim_out])  # pylint: disable=no-explicit-dtype
+            self.bias_atom_e = np.zeros(
+                [self.ntypes, net_dim_out], dtype=GLOBAL_NP_FLOAT_PRECISION
+            )
         else:
             assert bias_atom_e.shape == (self.ntypes, net_dim_out)
-            self.bias_atom_e = bias_atom_e
+            self.bias_atom_e = bias_atom_e.astype(GLOBAL_NP_FLOAT_PRECISION)
         if self.numb_fparam > 0:
-            self.fparam_avg = np.zeros(self.numb_fparam)  # pylint: disable=no-explicit-dtype
-            self.fparam_inv_std = np.ones(self.numb_fparam)  # pylint: disable=no-explicit-dtype
+            self.fparam_avg = np.zeros(self.numb_fparam, dtype=self.prec)
+            self.fparam_inv_std = np.ones(self.numb_fparam, dtype=self.prec)
         else:
             self.fparam_avg, self.fparam_inv_std = None, None
         if self.numb_aparam > 0:
-            self.aparam_avg = np.zeros(self.numb_aparam)  # pylint: disable=no-explicit-dtype
-            self.aparam_inv_std = np.ones(self.numb_aparam)  # pylint: disable=no-explicit-dtype
+            self.aparam_avg = np.zeros(self.numb_aparam, dtype=self.prec)
+            self.aparam_inv_std = np.ones(self.numb_aparam, dtype=self.prec)
         else:
             self.aparam_avg, self.aparam_inv_std = None, None
         # init networks
@@ -410,7 +422,9 @@ def _call_common(
 
         # calcualte the prediction
         if not self.mixed_types:
-            outs = xp.zeros([nf, nloc, net_dim_out])  # pylint: disable=no-explicit-dtype
+            outs = xp.zeros(
+                [nf, nloc, net_dim_out], dtype=get_xp_precision(xp, self.precision)
+            )
             for type_i in range(self.ntypes):
                 mask = xp.tile(
                     xp.reshape((atype == type_i), [nf, nloc, 1]), (1, 1, net_dim_out)
@@ -436,4 +450,4 @@ def _call_common(
         exclude_mask = self.emask.build_type_exclude_mask(atype)
         # nf x nloc x nod
         outs = outs * xp.astype(exclude_mask[:, :, None], outs.dtype)
-        return {self.var_name: outs}
+        return {self.var_name: xp.astype(outs, get_xp_precision(xp, "global"))}

deepmd/dpmodel/fitting/polarizability_fitting.py

@@ -303,7 +303,7 @@ def call(
             bias = self.constant_matrix[atype]
             # (nframes, nloc, 1)
             bias = np.expand_dims(bias, axis=-1) * self.scale[atype]
-            eye = np.eye(3)  # pylint: disable=no-explicit-dtype
+            eye = np.eye(3, dtype=descriptor.dtype)
             eye = np.tile(eye, (nframes, nloc, 1, 1))
             # (nframes, nloc, 3, 3)
             bias = np.expand_dims(bias, axis=-1) * eye

deepmd/dpmodel/infer/deep_eval.py

@@ -24,6 +24,9 @@
 from deepmd.dpmodel.utils.serialization import (
     load_dp_model,
 )
+from deepmd.env import (
+    GLOBAL_NP_FLOAT_PRECISION,
+)
 from deepmd.infer.deep_dipole import (
     DeepDipole,
 )
@@ -340,12 +343,12 @@ def _eval_model(
                 if batch_output[dp_name] is not None:
                     out = batch_output[dp_name].reshape(shape)
                 else:
-                    out = np.full(shape, np.nan)  # pylint: disable=no-explicit-dtype
+                    out = np.full(shape, np.nan, dtype=GLOBAL_NP_FLOAT_PRECISION)
                 results.append(out)
             else:
                 shape = self._get_output_shape(odef, nframes, natoms)
                 results.append(
-                    np.full(np.abs(shape), np.nan)  # pylint: disable=no-explicit-dtype
+                    np.full(np.abs(shape), np.nan, dtype=GLOBAL_NP_FLOAT_PRECISION)
                 )  # this is kinda hacky
         return tuple(results)
 

deepmd/dpmodel/utils/nlist.py

@@ -131,7 +131,7 @@ def build_neighbor_list(
         nlist = nlist[:, :, :nsel]
     else:
         rr = xp.concatenate(
-            [rr, xp.ones([batch_size, nloc, nsel - nnei]) + rcut],  # pylint: disable=no-explicit-dtype
+            [rr, xp.ones([batch_size, nloc, nsel - nnei], dtype=rr.dtype) + rcut],
             axis=-1,
         )
         nlist = xp.concatenate(
@@ -277,7 +277,8 @@ def extend_coord_with_ghosts(
     """
     xp = array_api_compat.array_namespace(coord, atype)
     nf, nloc = atype.shape
-    aidx = xp.tile(xp.arange(nloc)[xp.newaxis, :], (nf, 1))  # pylint: disable=no-explicit-dtype
+    # int64 for index
+    aidx = xp.tile(xp.arange(nloc, dtype=xp.int64)[xp.newaxis, :], (nf, 1))
     if cell is None:
         nall = nloc
         extend_coord = coord
@@ -289,9 +290,9 @@ def extend_coord_with_ghosts(
         to_face = to_face_distance(cell)
         nbuff = xp.astype(xp.ceil(rcut / to_face), xp.int64)
         nbuff = xp.max(nbuff, axis=0)
-        xi = xp.arange(-int(nbuff[0]), int(nbuff[0]) + 1, 1)  # pylint: disable=no-explicit-dtype
-        yi = xp.arange(-int(nbuff[1]), int(nbuff[1]) + 1, 1)  # pylint: disable=no-explicit-dtype
-        zi = xp.arange(-int(nbuff[2]), int(nbuff[2]) + 1, 1)  # pylint: disable=no-explicit-dtype
+        xi = xp.arange(-int(nbuff[0]), int(nbuff[0]) + 1, 1, dtype=xp.int64)
+        yi = xp.arange(-int(nbuff[1]), int(nbuff[1]) + 1, 1, dtype=xp.int64)
+        zi = xp.arange(-int(nbuff[2]), int(nbuff[2]) + 1, 1, dtype=xp.int64)
         xyz = xp.linalg.outer(xi, xp.asarray([1, 0, 0]))[:, xp.newaxis, xp.newaxis, :]
         xyz = (
             xyz

deepmd/infer/model_devi.py

@@ -328,7 +328,7 @@ def calc_model_devi(
     forces = np.array(forces)
     virials = np.array(virials)
 
-    devi = [np.arange(coord.shape[0]) * frequency]  # pylint: disable=no-explicit-dtype
+    devi = [np.arange(coord.shape[0], dtype=np.int64) * frequency]
     if real_data is None:
         devi += list(calc_model_devi_v(virials, relative=relative_v))
         devi_f = list(calc_model_devi_f(forces, relative=relative, atomic=atomic))
@@ -502,7 +502,7 @@ def make_model_devi(
             nframes_tot += coord.shape[0]
             devis.append(devi)
         devis = np.vstack(devis)
-        devis[:, 0] = np.arange(nframes_tot) * frequency  # pylint: disable=no-explicit-dtype
+        devis[:, 0] = np.arange(nframes_tot, dtype=np.int64) * frequency
         write_model_devi_out(devis, output, header=system, atomic=atomic)
         devis_coll.append(devis)
     return devis_coll

deepmd/pt/infer/deep_eval.py

@@ -396,6 +396,7 @@ def _eval_model(
         request_defs: list[OutputVariableDef],
     ):
         model = self.dp.to(DEVICE)
+        prec = NP_PRECISION_DICT[RESERVED_PRECISON_DICT[GLOBAL_PT_FLOAT_PRECISION]]
 
         nframes = coords.shape[0]
         if len(atom_types.shape) == 1:
@@ -405,9 +406,7 @@ def _eval_model(
             natoms = len(atom_types[0])
 
         coord_input = torch.tensor(
-            coords.reshape([nframes, natoms, 3]).astype(
-                NP_PRECISION_DICT[RESERVED_PRECISON_DICT[GLOBAL_PT_FLOAT_PRECISION]]
-            ),
+            coords.reshape([nframes, natoms, 3]).astype(prec),
             dtype=GLOBAL_PT_FLOAT_PRECISION,
             device=DEVICE,
         )
@@ -418,9 +417,7 @@ def _eval_model(
         )
         if cells is not None:
             box_input = torch.tensor(
-                cells.reshape([nframes, 3, 3]).astype(
-                    NP_PRECISION_DICT[RESERVED_PRECISON_DICT[GLOBAL_PT_FLOAT_PRECISION]]
-                ),
+                cells.reshape([nframes, 3, 3]).astype(prec),
                 dtype=GLOBAL_PT_FLOAT_PRECISION,
                 device=DEVICE,
             )
@@ -462,7 +459,7 @@ def _eval_model(
             else:
                 shape = self._get_output_shape(odef, nframes, natoms)
                 results.append(
-                    np.full(np.abs(shape), np.nan)  # pylint: disable=no-explicit-dtype
+                    np.full(np.abs(shape), np.nan, dtype=prec)
                 )  # this is kinda hacky
         return tuple(results)
 
@@ -542,7 +539,13 @@ def _eval_model_spin(
             else:
                 shape = self._get_output_shape(odef, nframes, natoms)
                 results.append(
-                    np.full(np.abs(shape), np.nan)  # pylint: disable=no-explicit-dtype
+                    np.full(
+                        np.abs(shape),
+                        np.nan,
+                        dtype=NP_PRECISION_DICT[
+                            RESERVED_PRECISON_DICT[GLOBAL_PT_FLOAT_PRECISION]
+                        ],
+                    )
                 )  # this is kinda hacky
         return tuple(results)
 

deepmd/pt/model/atomic_model/linear_atomic_model.py

@@ -90,7 +90,9 @@ def __init__(
         self.rcuts = torch.tensor(
             self.get_model_rcuts(), dtype=torch.float64, device=env.DEVICE
         )
-        self.nsels = torch.tensor(self.get_model_nsels(), device=env.DEVICE)  # pylint: disable=no-explicit-dtype
+        self.nsels = torch.tensor(
+            self.get_model_nsels(), device=env.DEVICE, dtype=torch.int32
+        )
 
         if isinstance(weights, str):
             assert weights in ["sum", "mean"]
@@ -299,8 +301,11 @@ def remap_atype(ori_map: list[str], new_map: list[str]) -> torch.Tensor:
         """
         type_2_idx = {atp: idx for idx, atp in enumerate(ori_map)}
         # this maps the atype in the new map to the original map
-        mapping = torch.tensor(  # pylint: disable=no-explicit-dtype
-            [type_2_idx[new_map[idx]] for idx in range(len(new_map))], device=env.DEVICE
+        # int32 should be enough for number of atom types.
+        mapping = torch.tensor(
+            [type_2_idx[new_map[idx]] for idx in range(len(new_map))],
+            device=env.DEVICE,
+            dtype=torch.int32,
         )
         return mapping
 

deepmd/pt/model/atomic_model/pairtab_atomic_model.py

@@ -269,9 +269,11 @@ def forward_atomic(
         # i_type : (nframes, nloc), this is atype.
         # j_type : (nframes, nloc, nnei)
         j_type = extended_atype[
-            torch.arange(extended_atype.size(0), device=extended_coord.device)[  # pylint: disable=no-explicit-dtype
-                :, None, None
-            ],
+            torch.arange(
+                extended_atype.size(0),
+                device=extended_coord.device,
+                dtype=torch.int64,
+            )[:, None, None],
             masked_nlist,
         ]
 

deepmd/pt/model/descriptor/descriptor.py

@@ -145,8 +145,16 @@ def share_params(self, base_class, shared_level, resume=False):
                         base_env.stats[kk] += self.get_stats()[kk]
                     mean, stddev = base_env()
                     if not base_class.set_davg_zero:
-                        base_class.mean.copy_(torch.tensor(mean, device=env.DEVICE))  # pylint: disable=no-explicit-dtype
-                    base_class.stddev.copy_(torch.tensor(stddev, device=env.DEVICE))  # pylint: disable=no-explicit-dtype
+                        base_class.mean.copy_(
+                            torch.tensor(
+                                mean, device=env.DEVICE, dtype=base_class.mean.dtype
+                            )
+                        )
+                    base_class.stddev.copy_(
+                        torch.tensor(
+                            stddev, device=env.DEVICE, dtype=base_class.stddev.dtype
+                        )
+                    )
                 # must share, even if not do stat
                 self.mean = base_class.mean
                 self.stddev = base_class.stddev