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)

source/tests/jax/test_fitting.py

@@ -0,0 +1,117 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+import unittest
+
+import array_api_compat
+import numpy as np
+from jax import (
+    tree_util,
+)
+
+from deepmd.jax.env import (
+    jnp,
+    nnx,
+)
+from deepmd.jax.fitting.fitting import (
+    EnergyFittingNet,
+)
+
+
+class TestJAXFitting(unittest.TestCase):
+    def test_runtime_buffers_support_nnx_jit_and_grad_tracing(self) -> None:
+        """Portable fitting buffers must remain usable inside JAX tracing."""
+        fitting = EnergyFittingNet(
+            ntypes=2,
+            dim_descrpt=3,
+            neuron=[5, 5],
+            numb_fparam=2,
+            numb_aparam=1,
+            dim_case_embd=2,
+            default_fparam=[0.5, -0.25],
+            precision="float32",
+            mixed_types=True,
+            seed=20260724,
+        )
+        # Non-default values ensure every portable buffer contributes to the
+        # traced result instead of being optimized away as an all-zero no-op.
+        buffer_values = {
+            "bias_atom_e": np.array([[1.5], [-0.75]], dtype=np.float32),
+            "fparam_avg": np.array([0.1, -0.2], dtype=np.float32),
+            "fparam_inv_std": np.array([2.0, 0.5], dtype=np.float32),
+            "aparam_avg": np.array([0.25], dtype=np.float32),
+            "aparam_inv_std": np.array([4.0], dtype=np.float32),
+            "case_embd": np.array([0.3, -0.6], dtype=np.float32),
+            "default_fparam_tensor": np.array([0.5, -0.25], dtype=np.float32),
+        }
+        for name, value in buffer_values.items():
+            fitting[name] = value
+
+        descriptor = jnp.asarray(
+            [[[0.2, -0.1, 0.4], [0.5, 0.3, -0.2], [-0.4, 0.7, 0.1]]],
+            dtype=jnp.float32,
+        )
+        atype = jnp.asarray([[0, 1, 0]], dtype=jnp.int32)
+        aparam = jnp.asarray([[[0.5], [0.0], [1.0]]], dtype=jnp.float32)
+
+        @nnx.jit
+        def forward(model, descriptor, atype, aparam):
+            return model(descriptor, atype, aparam=aparam)["energy"]
+
+        @nnx.jit
+        def grad(model, descriptor, atype, aparam):
+            def loss(model):
+                return jnp.sum(model(descriptor, atype, aparam=aparam)["energy"])
+
+            return nnx.grad(loss)(model)
+
+        buffer_snapshots = {
+            name: (
+                fitting[name],
+                type(fitting[name]),
+                array_api_compat.array_namespace(fitting[name]),
+                np.asarray(fitting[name]).copy(),
+            )
+            for name in buffer_values
+        }
+
+        out = forward(fitting, descriptor, atype, aparam)
+        grad_state = grad(fitting, descriptor, atype, aparam)
+        grad_leaves = [np.asarray(leaf) for leaf in tree_util.tree_leaves(grad_state)]
+        self.assertEqual(tuple(out.shape), (1, 3, 1))
+        self.assertFalse(np.any(np.isnan(np.asarray(out))))
+        self.assertTrue(grad_leaves)
+        self.assertTrue(all(np.all(np.isfinite(leaf)) for leaf in grad_leaves))
+
+        # Tracing may read portable buffers, but it must not replace their NNX
+        # variables, change their array namespace, or mutate their values.
+        for name, (
+            original,
+            original_type,
+            namespace,
+            values,
+        ) in buffer_snapshots.items():
+            current = fitting[name]
+            self.assertIs(current, original)
+            self.assertIs(type(current), original_type)
+            self.assertIs(array_api_compat.array_namespace(current), namespace)
+            np.testing.assert_array_equal(np.asarray(current), values)
+
+        neutral_values = {
+            "bias_atom_e": np.zeros((2, 1), dtype=np.float32),
+            "fparam_avg": np.zeros(2, dtype=np.float32),
+            "fparam_inv_std": np.ones(2, dtype=np.float32),
+            "aparam_avg": np.zeros(1, dtype=np.float32),
+            "aparam_inv_std": np.ones(1, dtype=np.float32),
+            "case_embd": np.zeros(2, dtype=np.float32),
+            "default_fparam_tensor": np.zeros(2, dtype=np.float32),
+        }
+        baseline = np.asarray(out)
+        for name, neutral in neutral_values.items():
+            with self.subTest(buffer=name):
+                fitting[name] = neutral
+                changed = np.asarray(forward(fitting, descriptor, atype, aparam))
+                self.assertFalse(np.allclose(changed, baseline))
+                fitting[name] = buffer_values[name]
+
+
+if __name__ == "__main__":
+    unittest.main()