fix(dpmodel): materialize fitting buffers on active backend

deepmd/dpmodel/fitting/general_fitting.py

@@ -666,8 +666,16 @@ def _call_common(
         if self.numb_fparam > 0 and fparam is None:
             # use default fparam
             assert self.default_fparam_tensor is not None
+            # Fitting statistics remain NumPy for portable serialization.
+            # Materialize constants next to the runtime descriptor instead of
+            # asking a backend namespace to reshape a foreign array directly.
+            default_fparam_tensor = xp.asarray(
+                self.default_fparam_tensor,
+                dtype=descriptor.dtype,
+                device=array_api_compat.device(descriptor),
+            )
             fparam = xp.tile(
-                xp.reshape(self.default_fparam_tensor, (1, self.numb_fparam)), (nf, 1)
+                xp.reshape(default_fparam_tensor, (1, self.numb_fparam)), (nf, 1)
             )
 
         # check fparam dim, concate to input descriptor
@@ -680,7 +688,18 @@ def _call_common(
                     f"input fparam: cannot reshape {fparam.shape} "
                     f"into ({nf}, {self.numb_fparam})."
                 ) from e
-            fparam = (fparam - self.fparam_avg[...]) * self.fparam_inv_std[...]
+            fparam_device = array_api_compat.device(fparam)
+            fparam_avg = xp.asarray(
+                self.fparam_avg,
+                dtype=fparam.dtype,
+                device=fparam_device,
+            )
+            fparam_inv_std = xp.asarray(
+                self.fparam_inv_std,
+                dtype=fparam.dtype,
+                device=fparam_device,
+            )
+            fparam = (fparam - fparam_avg) * fparam_inv_std
             fparam = xp.tile(
                 xp.reshape(fparam, (nf, 1, self.numb_fparam)), (1, nloc, 1)
             )
@@ -703,7 +722,18 @@ def _call_common(
                     f"input aparam: cannot reshape {aparam.shape} "
                     f"into ({nf}, {nloc}, {self.numb_aparam})."
                 ) from e
-            aparam = (aparam - self.aparam_avg[...]) * self.aparam_inv_std[...]
+            aparam_device = array_api_compat.device(aparam)
+            aparam_avg = xp.asarray(
+                self.aparam_avg,
+                dtype=aparam.dtype,
+                device=aparam_device,
+            )
+            aparam_inv_std = xp.asarray(
+                self.aparam_inv_std,
+                dtype=aparam.dtype,
+                device=aparam_device,
+            )
+            aparam = (aparam - aparam_avg) * aparam_inv_std
             xx = xp.concat(
                 [xx, aparam],
                 axis=-1,
@@ -716,9 +746,12 @@ def _call_common(
 
         if self.dim_case_embd > 0:
             assert self.case_embd is not None
-            case_embd = xp.tile(
-                xp.reshape(self.case_embd[...], (1, 1, -1)), (nf, nloc, 1)
+            case_embd_buffer = xp.asarray(
+                self.case_embd,
+                dtype=descriptor.dtype,
+                device=array_api_compat.device(descriptor),
             )
+            case_embd = xp.tile(xp.reshape(case_embd_buffer, (1, 1, -1)), (nf, nloc, 1))
             xx = xp.concat(
                 [xx, case_embd],
                 axis=-1,
@@ -773,9 +806,14 @@ def _call_common(
                 outs -= self.nets[()](xx_zeros)
             if self.eval_return_middle_output and len(self.neuron) > 0:
                 middle_outs = self.nets[()].call_until_last(xx)
+        bias_atom_e = xp.asarray(
+            self.bias_atom_e,
+            dtype=outs.dtype,
+            device=array_api_compat.device(outs),
+        )
         outs += xp.reshape(
             xp.take(
-                xp.astype(self.bias_atom_e[...], outs.dtype),
+                bias_atom_e,
                 xp.reshape(atype, (-1,)),
                 axis=0,
             ),

deepmd/dpmodel/fitting/polarizability_fitting.py

@@ -330,8 +330,16 @@ def call(
         results = self._call_common(descriptor, atype, gr, g2, h2, fparam, aparam)
         out = results.pop(self.var_name)
         # out = out * self.scale[atype, ...]
+        # Scale and diagonal shifts are serialized as NumPy arrays. Keep that
+        # portable representation, but gather them on the prediction backend
+        # so direct dpmodel use does not depend on eager wrapper conversion.
+        scale = xp.asarray(
+            self.scale,
+            dtype=out.dtype,
+            device=array_api_compat.device(out),
+        )
         scale_atype = xp.reshape(
-            xp.take(xp.astype(self.scale, out.dtype), xp.reshape(atype, (-1,)), axis=0),
+            xp.take(scale, xp.reshape(atype, (-1,)), axis=0),
             (*atype.shape, 1),
         )
         out = out * scale_atype
@@ -354,9 +362,14 @@ def call(
         out = xp.reshape(out, (nframes, nloc, 3, 3))
         if self.shift_diag:
             # bias = self.constant_matrix[atype]
+            constant_matrix = xp.asarray(
+                self.constant_matrix,
+                dtype=out.dtype,
+                device=array_api_compat.device(out),
+            )
             bias = xp.reshape(
                 xp.take(
-                    xp.astype(self.constant_matrix, out.dtype),
+                    constant_matrix,
                     xp.reshape(atype, (-1,)),
                     axis=0,
                 ),

source/tests/common/dpmodel/test_fitting_invar_fitting.py

@@ -4,6 +4,11 @@
 
 import numpy as np
 
+try:
+    import torch
+except ImportError:
+    torch = None
+
 from deepmd.dpmodel.descriptor import (
     DescrptSeA,
 )
@@ -177,3 +182,48 @@ def test_get_set(self) -> None:
         ]:
             ifn0[ii] = foo
             np.testing.assert_allclose(foo, ifn0[ii])
+
+    @unittest.skipIf(torch is None, "PyTorch is not installed")
+    def test_runtime_buffers_follow_torch_descriptor(self) -> None:
+        """Portable fitting buffers must materialize on the active backend."""
+        fitting = InvarFitting(
+            "energy",
+            ntypes=2,
+            dim_descrpt=3,
+            dim_out=1,
+            neuron=[5, 5],
+            numb_fparam=2,
+            numb_aparam=1,
+            dim_case_embd=2,
+            default_fparam=[0.5, -0.25],
+            precision="float64",
+            mixed_types=True,
+            seed=20260717,
+        )
+        fitting.bias_atom_e[:] = [[1.5], [-0.75]]
+        fitting.fparam_avg[:] = [0.1, -0.2]
+        fitting.fparam_inv_std[:] = [2.0, 0.5]
+        fitting.aparam_avg[:] = [0.25]
+        fitting.aparam_inv_std[:] = [4.0]
+        fitting.case_embd[:] = [0.3, -0.6]
+
+        descriptor = np.array(
+            [[[0.2, -0.1, 0.4], [0.5, 0.3, -0.2], [-0.4, 0.7, 0.1]]],
+            dtype=np.float64,
+        )
+        atype = np.array([[0, 1, 0]], dtype=np.int64)
+        aparam = np.array([[[0.5], [0.0], [1.0]]], dtype=np.float64)
+        expected = fitting(descriptor, atype, aparam=aparam)["energy"]
+
+        # Backend wrappers eagerly convert these NumPy attributes and would
+        # hide the generic dpmodel boundary, so convert only runtime inputs.
+        result = fitting(
+            torch.as_tensor(descriptor),
+            torch.as_tensor(atype),
+            aparam=torch.as_tensor(aparam),
+        )["energy"]
+
+        self.assertIsInstance(result, torch.Tensor)
+        self.assertEqual(result.dtype, torch.float64)
+        self.assertEqual(result.device.type, "cpu")
+        np.testing.assert_allclose(result.detach().cpu().numpy(), expected)

source/tests/common/dpmodel/test_fitting_polar_backend.py

@@ -0,0 +1,59 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+"""Backend-boundary regressions for dpmodel polarizability fitting."""
+
+import unittest
+
+import numpy as np
+
+try:
+    import torch
+except ImportError:
+    torch = None
+
+from deepmd.dpmodel.fitting.polarizability_fitting import (
+    PolarFitting,
+)
+
+
+class TestPolarFittingBackendBuffers(unittest.TestCase):
+    """Keep portable polar buffers compatible with backend-native inputs."""
+
+    @unittest.skipIf(torch is None, "PyTorch is not installed")
+    def test_scale_and_shift_follow_torch_prediction(self) -> None:
+        fitting = PolarFitting(
+            ntypes=2,
+            dim_descrpt=3,
+            embedding_width=2,
+            neuron=[5, 5],
+            precision="float64",
+            mixed_types=True,
+            fit_diag=True,
+            scale=[1.5, 0.25],
+            shift_diag=True,
+            seed=20260717,
+        )
+        fitting.constant_matrix[:] = [0.75, -0.5]
+
+        descriptor = np.array([[[0.2, -0.1, 0.4], [0.5, 0.3, -0.2]]], dtype=np.float64)
+        atype = np.array([[0, 1]], dtype=np.int64)
+        rotation = np.array(
+            [
+                [
+                    [[1.0, 0.0, 0.0], [0.0, 1.0, 0.0]],
+                    [[0.0, 1.0, 0.0], [0.0, 0.0, 1.0]],
+                ]
+            ],
+            dtype=np.float64,
+        )
+        expected = fitting(descriptor, atype, gr=rotation)["polarizability"]
+
+        result = fitting(
+            torch.as_tensor(descriptor),
+            torch.as_tensor(atype),
+            gr=torch.as_tensor(rotation),
+        )["polarizability"]
+
+        self.assertIsInstance(result, torch.Tensor)
+        self.assertEqual(result.dtype, torch.float64)
+        self.assertEqual(result.device.type, "cpu")
+        np.testing.assert_allclose(result.detach().cpu().numpy(), expected)