diff --git a/deepmd/dpmodel/fitting/general_fitting.py b/deepmd/dpmodel/fitting/general_fitting.py index 4734a6be66..6abc3ebf83 100644 --- a/deepmd/dpmodel/fitting/general_fitting.py +++ b/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, ), diff --git a/deepmd/dpmodel/fitting/polarizability_fitting.py b/deepmd/dpmodel/fitting/polarizability_fitting.py index e55a29c36a..f4e504a7aa 100644 --- a/deepmd/dpmodel/fitting/polarizability_fitting.py +++ b/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, ), diff --git a/source/tests/common/dpmodel/test_fitting_invar_fitting.py b/source/tests/common/dpmodel/test_fitting_invar_fitting.py index c93e46f66b..0590feae07 100644 --- a/source/tests/common/dpmodel/test_fitting_invar_fitting.py +++ b/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) diff --git a/source/tests/common/dpmodel/test_fitting_polar_backend.py b/source/tests/common/dpmodel/test_fitting_polar_backend.py new file mode 100644 index 0000000000..9de8357f0d --- /dev/null +++ b/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)