fix(dpmodel): handle cross-namespace device arrays

Author: OutisLi

deepmd/dpmodel/array_api.py

@@ -11,6 +11,10 @@
     Version,
 )
 
+from deepmd.dpmodel.common import (
+    to_numpy_array,
+)
+
 # Type alias for array_api compatible arrays
 Array = np.ndarray | Any  # Any to support JAX, PyTorch, etc. arrays
 
@@ -27,22 +31,28 @@ def xp_asarray_nodetach(
     ``torch.asarray`` detaches its input from the autograd graph, so calling
     ``xp.asarray`` on a weight attribute that is already a backend tensor
     (e.g. a ``torch.nn.Parameter`` registered by the pt_expt backend)
-    silently breaks gradient flow to that weight.  This helper converts
-    genuine non-backend data (numpy arrays, python scalars/lists) via
-    ``xp.asarray``; backend tensors are returned as-is, with an optional
-    differentiable dtype cast via ``xp.astype``.
-
-    The ``device`` argument only applies to the conversion path: backend
-    tensors are assumed to already live on the working device (they are
-    created together with the inputs).
+    silently breaks gradient flow to that weight. Backend tensors already in
+    ``xp`` are therefore returned as-is, with an optional differentiable dtype
+    cast via ``xp.astype``.
+
+    An array from another namespace cannot retain its autograd graph. It is
+    converted through NumPy before entering ``xp``; this also performs the
+    required device-to-host copy when a CUDA-backed model constant is consumed
+    by a NumPy statistics path.
+
+    The ``device`` argument only applies to the conversion path. Arrays already
+    in ``xp`` are assumed to live on the working device because model buffers
+    and inputs are moved together.
     """
-    if isinstance(obj, np.ndarray) or not array_api_compat.is_array_api_obj(obj):
-        if dtype is None:
-            return xp.asarray(obj, device=device)
-        return xp.asarray(obj, dtype=dtype, device=device)
-    if dtype is not None and obj.dtype != dtype:
-        obj = xp.astype(obj, dtype)
-    return obj
+    if array_api_compat.is_array_api_obj(obj):
+        if array_api_compat.array_namespace(obj) is xp:
+            if dtype is not None and obj.dtype != dtype:
+                obj = xp.astype(obj, dtype)
+            return obj
+        obj = to_numpy_array(obj)
+    if dtype is None:
+        return xp.asarray(obj, device=device)
+    return xp.asarray(obj, dtype=dtype, device=device)
 
 
 # array api adds take_along_axis in https://github.com/data-apis/array-api/pull/816

deepmd/dpmodel/descriptor/dpa4.py

@@ -2186,7 +2186,7 @@ def _canonicalize_charge_spin(
                 raise ValueError("`charge_spin` is required for this SeZM descriptor.")
             charge_spin = xp.reshape(
                 xp_asarray_nodetach(
-                    xp, np.asarray(self.default_chg_spin), dtype=dtype, device=device
+                    xp, self.default_chg_spin, dtype=dtype, device=device
                 ),
                 (1, 2),
             )

deepmd/dpmodel/model/spin_model.py

@@ -14,13 +14,13 @@
 
 from deepmd.dpmodel.array_api import (
     Array,
+    xp_asarray_nodetach,
 )
 from deepmd.dpmodel.atomic_model.dp_atomic_model import (
     DPAtomicModel,
 )
 from deepmd.dpmodel.common import (
     NativeOP,
-    to_numpy_array,
 )
 from deepmd.dpmodel.model.base_model import (
     BaseModel,
@@ -140,9 +140,11 @@ def __init__(
 
     def _to_xp(self, arr: Any, xp: Any, ref_arr: Any) -> Any:
         """Convert an array to the namespace and device of ``ref_arr``."""
-        if array_api_compat.is_numpy_namespace(xp):
-            arr = to_numpy_array(arr)
-        return xp.asarray(arr, device=array_api_compat.device(ref_arr))
+        return xp_asarray_nodetach(
+            xp,
+            arr,
+            device=array_api_compat.device(ref_arr),
+        )
 
     def _lookup_type_values(self, values: Any, atype: Array, ref_arr: Array) -> Array:
         """Gather per-type values while mapping virtual atom types to zero.

deepmd/dpmodel/utils/exclude_mask.py

@@ -5,6 +5,7 @@
 
 from deepmd.dpmodel.array_api import (
     Array,
+    xp_asarray_nodetach,
     xp_take_along_axis,
     xp_take_first_n,
 )
@@ -55,7 +56,11 @@ def build_type_exclude_mask(
         lead = atype.shape  # (nf, natom) dense | (N,) graph
         return xp.reshape(
             xp.take(
-                xp.asarray(self.type_mask[...], device=array_api_compat.device(atype)),
+                xp_asarray_nodetach(
+                    xp,
+                    self.type_mask[...],
+                    device=array_api_compat.device(atype),
+                ),
                 xp.reshape(atype, (-1,)),
                 axis=0,
             ),
@@ -151,7 +156,11 @@ def build_type_exclude_mask(
         type_ij_flat = xp.reshape(type_ij, (-1,))
         mask = xp.reshape(
             xp.take(
-                xp.asarray(self.type_mask[...], device=array_api_compat.device(nlist)),
+                xp_asarray_nodetach(
+                    xp,
+                    self.type_mask[...],
+                    device=array_api_compat.device(nlist),
+                ),
                 type_ij_flat,
             ),
             (nf, nloc, nnei),
@@ -185,7 +194,11 @@ def build_edge_exclude_mask(self, edge_index: Array, atype: Array) -> Array:
         dst_t = xp.take(atype, edge_index[1, :], axis=0)
         type_ij = dst_t * (self.ntypes + 1) + src_t
         return xp.take(
-            xp.asarray(self.type_mask[...], device=array_api_compat.device(atype)),
+            xp_asarray_nodetach(
+                xp,
+                self.type_mask[...],
+                device=array_api_compat.device(atype),
+            ),
             type_ij,
             axis=0,
         )

source/tests/consistent/test_array_api.py

@@ -1,11 +1,16 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
 import sys
 import unittest
+from unittest.mock import (
+    patch,
+)
 
+import array_api_compat
 import numpy as np
 
 from deepmd.dpmodel.array_api import (
     xp_add_at,
+    xp_asarray_nodetach,
     xp_bincount,
     xp_maximum_at,
     xp_scatter_sum,
@@ -56,6 +61,36 @@ def test_torch_parameter_requires_grad(self) -> None:
         self.assertTrue(param.requires_grad)
         self.assertEqual(param.device, DEVICE)
 
+    @unittest.skipUnless(INSTALLED_PT, "PyTorch is not installed")
+    def test_foreign_tensor_is_converted_to_numpy_namespace(self) -> None:
+        tensor = torch.tensor([1.0, 2.0], dtype=torch.float64, device=DEVICE)
+        numpy_namespace = array_api_compat.array_namespace(np.empty(0))
+
+        # A CUDA tensor rejects the direct NumPy protocol. Simulate that
+        # boundary on CPU so the device-to-host fallback is exercised on every
+        # platform.
+        with patch.object(
+            torch.Tensor,
+            "numpy",
+            side_effect=TypeError("direct conversion is unavailable"),
+        ):
+            converted = xp_asarray_nodetach(numpy_namespace, tensor)
+
+        self.assertIsInstance(converted, np.ndarray)
+        np.testing.assert_allclose(converted, np.array([1.0, 2.0]))
+
+    @unittest.skipUnless(INSTALLED_PT, "PyTorch is not installed")
+    def test_native_tensor_keeps_its_autograd_graph(self) -> None:
+        tensor = torch.nn.Parameter(
+            torch.tensor([1.0, 2.0], dtype=torch.float64, device=DEVICE)
+        )
+        torch_namespace = array_api_compat.array_namespace(tensor)
+
+        converted = xp_asarray_nodetach(torch_namespace, tensor)
+
+        self.assertIs(converted, tensor)
+        self.assertTrue(converted.requires_grad)
+
 
 class TestXpMaximumAtConsistent(unittest.TestCase):
     """Test maximum-at identities that differ between backend primitives."""

source/tests/pt_expt/descriptor/test_dpa4.py

@@ -107,6 +107,30 @@ def test_consistency(self, use_env_seed, use_mapping) -> None:
             err_msg=err_msg,
         )
 
+    def test_default_charge_spin_uses_model_namespace(self) -> None:
+        """A device buffer supplies the default without a NumPy round trip."""
+        dtype = PRECISION_DICT["float64"]
+        descriptor = make_descriptor(
+            self.nt,
+            self.sel_mix,
+            self.rcut,
+            add_chg_spin_ebd=True,
+            default_chg_spin=[0.5, -0.5],
+        ).to(self.device)
+        coord_ext = torch.tensor(self.coord_ext, dtype=dtype, device=self.device)
+        atype_ext = torch.tensor(self.atype_ext, dtype=int, device=self.device)
+        nlist = torch.tensor(self.nlist, dtype=int, device=self.device)
+
+        with mock.patch.object(
+            torch.Tensor,
+            "numpy",
+            side_effect=TypeError("direct conversion is unavailable"),
+        ):
+            output = descriptor(coord_ext, atype_ext, nlist)[0]
+
+        assert output.device == self.device
+        assert torch.isfinite(output).all()
+
     def test_train_and_eval_amp_switches_are_independent(self) -> None:
         """Training follows ``use_amp``, evaluation follows ``DP_AMP_INFER``.
 

source/tests/pt_expt/utils/test_exclusion_mask.py

@@ -1,5 +1,8 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
 import unittest
+from unittest.mock import (
+    patch,
+)
 
 import numpy as np
 import torch
@@ -38,6 +41,13 @@ def test_build_type_exclude_mask(self) -> None:
         des = AtomExcludeMask(nt, exclude_types=exclude_types)
         mask = des.build_type_exclude_mask(torch.as_tensor(atype, device=env.DEVICE))
         np.testing.assert_equal(mask.detach().cpu().numpy(), expected_mask)
+        with patch.object(
+            torch.Tensor,
+            "numpy",
+            side_effect=TypeError("direct conversion is unavailable"),
+        ):
+            numpy_mask = des.build_type_exclude_mask(atype)
+        np.testing.assert_equal(numpy_mask, expected_mask)
 
     def test_type_mask_is_buffer(self) -> None:
         des = AtomExcludeMask(3, exclude_types=[0])
@@ -66,6 +76,29 @@ def test_build_type_exclude_mask(self) -> None:
             torch.as_tensor(self.atype_ext, device=env.DEVICE),
         )
         np.testing.assert_equal(mask.detach().cpu().numpy(), expected_mask)
+        with patch.object(
+            torch.Tensor,
+            "numpy",
+            side_effect=TypeError("direct conversion is unavailable"),
+        ):
+            numpy_mask = des.build_type_exclude_mask(self.nlist, self.atype_ext)
+        np.testing.assert_equal(
+            numpy_mask,
+            expected_mask,
+        )
+
+    def test_build_edge_exclude_mask_with_numpy_inputs(self) -> None:
+        des = PairExcludeMask(self.nt, exclude_types=[[0, 1]])
+        edge_index = np.array([[0, 1, 2, 3], [1, 0, 3, 2]], dtype=np.int64)
+        atype = np.array([0, 1, 0, 0], dtype=np.int32)
+
+        with patch.object(
+            torch.Tensor,
+            "numpy",
+            side_effect=TypeError("direct conversion is unavailable"),
+        ):
+            numpy_mask = des.build_edge_exclude_mask(edge_index, atype)
+        np.testing.assert_equal(numpy_mask, np.array([0, 0, 1, 1], dtype=np.int32))
 
     def test_type_mask_is_buffer(self) -> None:
         des = PairExcludeMask(self.nt, exclude_types=[[0, 1]])