diff --git a/deepmd/pd/model/network/utils.py b/deepmd/pd/model/network/utils.py index cea832e11a..6b158a59e0 100644 --- a/deepmd/pd/model/network/utils.py +++ b/deepmd/pd/model/network/utils.py @@ -39,8 +39,10 @@ def aggregate( else: bin_count = None - # make sure this operation is done on the same device of data and owners - output = paddle.zeros([num_owner, data.shape[1]]) + # preserve the input dtype instead of falling back to the default float; + # paddle.zeros creates the tensor on the current global device (it has no + # place argument), which matches data/owners under deepmd's global device. + output = paddle.zeros([num_owner, data.shape[1]], dtype=data.dtype) output = output.index_add_(owners, 0, data.astype(output.dtype)) if average: assert bin_count is not None diff --git a/source/tests/pd/model/test_aggregate.py b/source/tests/pd/model/test_aggregate.py new file mode 100644 index 0000000000..cea05bc20c --- /dev/null +++ b/source/tests/pd/model/test_aggregate.py @@ -0,0 +1,27 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Test that ``aggregate`` preserves the input dtype (float64 must not downcast).""" + +import unittest + +import numpy as np +import paddle + +from deepmd.pd.model.network.utils import ( + aggregate, +) + + +class TestAggregateDtype(unittest.TestCase): + def _data(self) -> tuple[paddle.Tensor, paddle.Tensor]: + # rows 0,1 -> owner 0 ; row 2 -> owner 1 + data = paddle.to_tensor([[1.0], [2.0], [3.0]], dtype="float64") + owners = paddle.to_tensor([0, 0, 1], dtype="int64") + return data, owners + + def test_sum_preserves_float64(self) -> None: + # exercises the shared ``output = paddle.zeros(..., dtype=data.dtype)`` + # allocation via the summation path + data, owners = self._data() + out = aggregate(data, owners, average=False, num_owner=2) + self.assertEqual(out.dtype, paddle.float64) + np.testing.assert_allclose(out.numpy().ravel(), [3.0, 3.0])