fix(pd): preserve dtype in Paddle RepFlow dynamic aggregation

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

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])