fix(dpmodel): keep output reductions on active backend

deepmd/dpmodel/model/edge_transform_output.py

@@ -21,6 +21,8 @@
 
 from deepmd.dpmodel.common import (
     GLOBAL_ENER_FLOAT_PRECISION,
+    RESERVED_PRECISION_DICT,
+    get_xp_precision,
 )
 from deepmd.dpmodel.output_def import (
     get_deriv_name,
@@ -121,6 +123,11 @@ def fit_output_to_model_output_graph(
 
     n_node = graph.n_node
     xp = array_api_compat.get_namespace(n_node)
+    # The configured energy precision is represented by a NumPy dtype class,
+    # which is not accepted as a dtype by every array namespace (notably Torch).
+    energy_dtype = get_xp_precision(
+        xp, RESERVED_PRECISION_DICT[GLOBAL_ENER_FLOAT_PRECISION]
+    )
     nf = n_node.shape[0]
     frame_id = frame_id_from_n_node(n_node)
     n_total = next(iter(fit_ret.values())).shape[0]
@@ -133,21 +140,21 @@ def fit_output_to_model_output_graph(
         if not vdef.reducible:
             continue
         kk_redu = get_reduce_name(kk)
-        vv_e = xp.astype(vv, GLOBAL_ENER_FLOAT_PRECISION)
+        vv_e = xp.astype(vv, energy_dtype)
         if owned is not None:
-            owned_e = xp.astype(owned, GLOBAL_ENER_FLOAT_PRECISION)
+            owned_e = xp.astype(owned, energy_dtype)
             vv_e = vv_e * xp.reshape(owned_e, (n_total, *([1] * (vv_e.ndim - 1))))
         redu = segment_sum(vv_e, frame_id, nf)  # (nf, *shape)
         if vdef.intensive:
             if mask is not None:
-                cnt_mask = xp.astype(mask, GLOBAL_ENER_FLOAT_PRECISION)
+                cnt_mask = xp.astype(mask, energy_dtype)
                 if owned is not None:
                     cnt_mask = cnt_mask * owned_e
                 cnt = segment_sum(cnt_mask, frame_id, nf)
             elif owned is not None:
                 cnt = segment_sum(owned_e, frame_id, nf)
             else:
-                cnt = xp.astype(n_node, GLOBAL_ENER_FLOAT_PRECISION)
+                cnt = xp.astype(n_node, energy_dtype)
             redu = redu / xp.reshape(cnt, (nf, *([1] * (redu.ndim - 1))))
         model_ret[kk_redu] = redu
         if vdef.r_differentiable:

deepmd/dpmodel/model/transform_output.py

@@ -2,14 +2,15 @@
 
 
 import array_api_compat
-import numpy as np
 
 from deepmd.dpmodel.array_api import (
     Array,
     xp_scatter_sum,
 )
 from deepmd.dpmodel.common import (
     GLOBAL_ENER_FLOAT_PRECISION,
+    RESERVED_PRECISION_DICT,
+    get_xp_precision,
 )
 from deepmd.dpmodel.output_def import (
     FittingOutputDef,
@@ -34,6 +35,12 @@ def fit_output_to_model_output(
 
     """
     xp = array_api_compat.get_namespace(coord_ext)
+    # GLOBAL_ENER_FLOAT_PRECISION is a NumPy dtype class.  Array namespaces such
+    # as Torch require their own dtype object even when the precision name is the
+    # same, so resolve it before casting any backend array.
+    energy_dtype = get_xp_precision(
+        xp, RESERVED_PRECISION_DICT[GLOBAL_ENER_FLOAT_PRECISION]
+    )
     model_ret = dict(fit_ret.items())
     for kk, vv in fit_ret.items():
         vdef = fit_output_def[kk]
@@ -45,16 +52,14 @@ def fit_output_to_model_output(
             if vdef.intensive:
                 if mask is not None:
                     model_ret[kk_redu] = xp.sum(
-                        vv.astype(GLOBAL_ENER_FLOAT_PRECISION), axis=atom_axis
-                    ) / np.sum(mask, axis=-1, keepdims=True)
+                        xp.astype(vv, energy_dtype), axis=atom_axis
+                    ) / xp.sum(xp.astype(mask, energy_dtype), axis=-1, keepdims=True)
                 else:
                     model_ret[kk_redu] = xp.mean(
-                        vv.astype(GLOBAL_ENER_FLOAT_PRECISION), axis=atom_axis
+                        xp.astype(vv, energy_dtype), axis=atom_axis
                     )
             else:
-                model_ret[kk_redu] = xp.sum(
-                    vv.astype(GLOBAL_ENER_FLOAT_PRECISION), axis=atom_axis
-                )
+                model_ret[kk_redu] = xp.sum(xp.astype(vv, energy_dtype), axis=atom_axis)
             if vdef.r_differentiable:
                 kk_derv_r, kk_derv_c = get_deriv_name(kk)
                 # name-holders

source/tests/common/dpmodel/test_transform_output.py

@@ -0,0 +1,143 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+"""Backend-preservation tests for fitting-output reductions."""
+
+import pytest
+
+from deepmd.dpmodel.model.edge_transform_output import (
+    fit_output_to_model_output_graph,
+)
+from deepmd.dpmodel.model.transform_output import (
+    fit_output_to_model_output,
+)
+from deepmd.dpmodel.output_def import (
+    FittingOutputDef,
+    OutputVariableDef,
+)
+from deepmd.dpmodel.utils.neighbor_graph import (
+    NeighborGraph,
+)
+
+torch = pytest.importorskip("torch")
+
+
+def _output_def(*, intensive: bool) -> FittingOutputDef:
+    """Build the smallest reducible fitting definition used by these tests."""
+    return FittingOutputDef(
+        [
+            OutputVariableDef(
+                name="energy",
+                shape=[1],
+                reducible=True,
+                r_differentiable=False,
+                c_differentiable=False,
+                intensive=intensive,
+            )
+        ]
+    )
+
+
+def test_dense_torch_extensive_reduction_stays_on_backend() -> None:
+    """Dense reduction must not rely on NumPy-style ``Tensor.astype``."""
+    atomic = torch.tensor([[[1.0], [2.0], [3.0]]], dtype=torch.float32)
+    coord = torch.zeros((1, 3, 3), dtype=torch.float32)
+
+    result = fit_output_to_model_output(
+        {"energy": atomic}, _output_def(intensive=False), coord
+    )["energy_redu"]
+
+    assert isinstance(result, torch.Tensor)
+    assert result.dtype is torch.float64
+    torch.testing.assert_close(result, torch.tensor([[6.0]], dtype=torch.float64))
+
+
+def test_dense_torch_mask_count_uses_energy_dtype() -> None:
+    """The intensive divisor must be reduced by Torch in energy precision."""
+    atomic = torch.tensor([[[1.0], [2.0], [3.0]]], dtype=torch.float32)
+    coord = torch.zeros((1, 3, 3), dtype=torch.float32)
+    mask = torch.tensor([[True, True, False]])
+
+    result = fit_output_to_model_output(
+        {"energy": atomic}, _output_def(intensive=True), coord, mask=mask
+    )["energy_redu"]
+
+    assert isinstance(result, torch.Tensor)
+    assert result.dtype is torch.float64
+    torch.testing.assert_close(result, torch.tensor([[3.0]], dtype=torch.float64))
+
+
+def test_dense_torch_no_mask_intensive_reduction_stays_on_backend() -> None:
+    """Exercise the intensive reduction branch without a real-atom mask."""
+    atomic = torch.tensor([[[1.0], [2.0], [4.0]]], dtype=torch.float32)
+    coord = torch.zeros((1, 3, 3), dtype=torch.float32)
+
+    result = fit_output_to_model_output(
+        {"energy": atomic}, _output_def(intensive=True), coord
+    )["energy_redu"]
+
+    assert isinstance(result, torch.Tensor)
+    assert result.dtype is torch.float64
+    torch.testing.assert_close(result, torch.tensor([[7.0 / 3.0]], dtype=torch.float64))
+
+
+def test_graph_torch_mask_count_uses_backend_dtype() -> None:
+    """Graph reductions must translate the NumPy precision to ``torch.dtype``."""
+    graph = NeighborGraph(
+        n_node=torch.tensor([2, 1], dtype=torch.int64),
+        edge_index=torch.empty((2, 0), dtype=torch.int64),
+        edge_vec=torch.empty((0, 3), dtype=torch.float32),
+        edge_mask=torch.empty((0,), dtype=torch.bool),
+    )
+    atomic = torch.tensor([[1.0], [2.0], [6.0]], dtype=torch.float32)
+    mask = torch.tensor([True, True, True])
+
+    result = fit_output_to_model_output_graph(
+        {"energy": atomic}, _output_def(intensive=True), graph, mask=mask
+    )["energy_redu"]
+
+    assert isinstance(result, torch.Tensor)
+    assert result.dtype is torch.float64
+    torch.testing.assert_close(
+        result, torch.tensor([[1.5], [6.0]], dtype=torch.float64)
+    )
+
+
+def test_graph_torch_no_mask_intensive_reduction_stays_on_backend() -> None:
+    """Exercise the graph intensive fallback using per-frame node counts."""
+    graph = NeighborGraph(
+        n_node=torch.tensor([2, 1], dtype=torch.int64),
+        edge_index=torch.empty((2, 0), dtype=torch.int64),
+        edge_vec=torch.empty((0, 3), dtype=torch.float32),
+        edge_mask=torch.empty((0,), dtype=torch.bool),
+    )
+    atomic = torch.tensor([[1.0], [3.0], [6.0]], dtype=torch.float32)
+
+    result = fit_output_to_model_output_graph(
+        {"energy": atomic}, _output_def(intensive=True), graph
+    )["energy_redu"]
+
+    assert isinstance(result, torch.Tensor)
+    assert result.dtype is torch.float64
+    torch.testing.assert_close(
+        result, torch.tensor([[2.0], [6.0]], dtype=torch.float64)
+    )
+
+
+def test_graph_torch_extensive_reduction_stays_on_backend() -> None:
+    """Exercise the graph extensive branch independently of intensive counts."""
+    graph = NeighborGraph(
+        n_node=torch.tensor([2, 1], dtype=torch.int64),
+        edge_index=torch.empty((2, 0), dtype=torch.int64),
+        edge_vec=torch.empty((0, 3), dtype=torch.float32),
+        edge_mask=torch.empty((0,), dtype=torch.bool),
+    )
+    atomic = torch.tensor([[1.0], [3.0], [6.0]], dtype=torch.float32)
+
+    result = fit_output_to_model_output_graph(
+        {"energy": atomic}, _output_def(intensive=False), graph
+    )["energy_redu"]
+
+    assert isinstance(result, torch.Tensor)
+    assert result.dtype is torch.float64
+    torch.testing.assert_close(
+        result, torch.tensor([[4.0], [6.0]], dtype=torch.float64)
+    )