From 5a6c07fe027b06726518bbdf5b9ef997acda011b Mon Sep 17 00:00:00 2001
From: Han Wang <wang_han@iapcm.ac.cn>
Date: Thu, 2 Jul 2026 14:39:18 +0800
Subject: [PATCH 1/2] fix(pd): preserve dtype in RepFlow dynamic aggregation

aggregate() allocated its output with paddle.zeros(...) without a dtype, so it
fell back to paddle's default float (float32) and then cast the input data to
that dtype before index_add_. For float64 RepFlow/DPA3 models this silently
downcast descriptor updates to float32.

Allocate the output with dtype=data.dtype so the input precision is preserved.
Add a test that the summation path keeps float64.

Fix #5688
---
 deepmd/pd/model/network/utils.py        |  5 +++--
 source/tests/pd/model/test_aggregate.py | 27 +++++++++++++++++++++++++
 2 files changed, 30 insertions(+), 2 deletions(-)
 create mode 100644 source/tests/pd/model/test_aggregate.py

diff --git a/deepmd/pd/model/network/utils.py b/deepmd/pd/model/network/utils.py
index cea832e11a..792d6967e8 100644
--- a/deepmd/pd/model/network/utils.py
+++ b/deepmd/pd/model/network/utils.py
@@ -39,8 +39,9 @@ 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]])
+    # make sure this operation is done on the same device of data and owners,
+    # and preserve the input dtype instead of falling back to the default float
+    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])

From 57328e5a995cd15e40cb4c93185aa16bb42ada8d Mon Sep 17 00:00:00 2001
From: Han Wang <wang_han@iapcm.ac.cn>
Date: Tue, 7 Jul 2026 11:34:32 +0800
Subject: [PATCH 2/2] docs(pd): correct device wording in aggregate comment

The prior comment claimed the zeros allocation is placed on the same device as data/owners, but paddle.zeros has no place argument and only follows the current global device. Reword to describe that accurately (it matches data/owners under deepmd's global-device convention) instead of implying an enforced per-tensor device guarantee. Addresses the review note on #5712; comment-only, no behavior change.
---
 deepmd/pd/model/network/utils.py | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/deepmd/pd/model/network/utils.py b/deepmd/pd/model/network/utils.py
index 792d6967e8..6b158a59e0 100644
--- a/deepmd/pd/model/network/utils.py
+++ b/deepmd/pd/model/network/utils.py
@@ -39,8 +39,9 @@ def aggregate(
     else:
         bin_count = None
 
-    # make sure this operation is done on the same device of data and owners,
-    # and preserve the input dtype instead of falling back to the default float
+    # 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: