fix similar issues in models

Author: Han Wang

deepmd/dpmodel/array_api.py

@@ -51,7 +51,10 @@ def xp_take_along_axis(arr: Array, indices: Array, axis: int) -> Array:
     else:
         indices = xp.reshape(indices, (0, 0))
 
-    offset = (xp.arange(indices.shape[0], dtype=indices.dtype) * m)[:, xp.newaxis]
+    dev = array_api_compat.device(indices)
+    offset = (xp.arange(indices.shape[0], dtype=indices.dtype, device=dev) * m)[
+        :, xp.newaxis
+    ]
     indices = xp.reshape(offset + indices, (-1,))
 
     out = xp.take(arr, indices)

deepmd/dpmodel/atomic_model/linear_atomic_model.py

@@ -337,9 +337,11 @@ def _compute_weight(
         xp = array_api_compat.array_namespace(extended_coord, extended_atype, nlists_)
         nmodels = len(self.models)
         nframes, nloc, _ = nlists_[0].shape
+        dev = array_api_compat.device(extended_coord)
         # the dtype of weights is the interface data type.
         return [
-            xp.ones((nframes, nloc, 1), dtype=GLOBAL_NP_FLOAT_PRECISION) / nmodels
+            xp.ones((nframes, nloc, 1), dtype=GLOBAL_NP_FLOAT_PRECISION, device=dev)
+            / nmodels
             for _ in range(nmodels)
         ]
 

deepmd/dpmodel/atomic_model/pairtab_atomic_model.py

@@ -220,8 +220,11 @@ def forward_atomic(
         )  # (nframes, nloc, nnei)
 
         # (nframes, nloc, nnei), index type is int64.
+        dev = array_api_compat.device(extended_atype)
         j_type = extended_atype[
-            xp.arange(extended_atype.shape[0], dtype=xp.int64)[:, None, None],
+            xp.arange(extended_atype.shape[0], dtype=xp.int64, device=dev)[
+                :, None, None
+            ],
             masked_nlist,
         ]
 
@@ -327,8 +330,11 @@ def _get_pairwise_dist(coords: Array, nlist: Array) -> Array:
             The pairwise distance between the atoms (nframes, nloc, nnei).
         """
         xp = array_api_compat.array_namespace(coords, nlist)
+        dev = array_api_compat.device(nlist)
         # index type is int64
-        batch_indices = xp.arange(nlist.shape[0], dtype=xp.int64)[:, None, None]
+        batch_indices = xp.arange(nlist.shape[0], dtype=xp.int64, device=dev)[
+            :, None, None
+        ]
         neighbor_atoms = coords[batch_indices, nlist]
         loc_atoms = coords[:, : nlist.shape[1], :]
         pairwise_dr = loc_atoms[:, :, None, :] - neighbor_atoms

deepmd/dpmodel/model/make_model.py

@@ -540,7 +540,12 @@ def _format_nlist(
                 ret = xp.concat(
                     [
                         nlist,
-                        -1 * xp.ones([n_nf, n_nloc, nnei - n_nnei], dtype=nlist.dtype),
+                        -1
+                        * xp.ones(
+                            [n_nf, n_nloc, nnei - n_nnei],
+                            dtype=nlist.dtype,
+                            device=array_api_compat.device(nlist),
+                        ),
                     ],
                     axis=-1,
                 )