fix(dpmodel): keep output reductions on active backend (#5789)

## Summary

- translate the configured NumPy energy precision into the dtype object
of the active array namespace;
- use namespace-native casts and reductions in the dense dpmodel output
transform;
- apply the same dtype correction to the ragged graph transform, which
otherwise passes a NumPy dtype class to Torch;
- add focused Torch regressions for extensive, masked-intensive, and
graph reductions.

## Why existing tests missed this

There were no direct tests exercising the backend-agnostic dpmodel
transform helpers with Torch arrays. Normal PyTorch and Paddle model
tests call their backend-specific transform implementations instead,
while the existing dpmodel graph tests use NumPy. As a result, neither
Tensor.astype nor the NumPy dtype passed into the Torch namespace was
reached.

The new tests invoke the generic dpmodel helpers directly and verify
that results remain Torch tensors in the configured energy precision.

## Validation

- new transform regressions plus existing ragged graph tests: 10 passed;
- additional reviewer probes passed for NumPy, Torch, JAX, and
array-api-strict;
- ruff format .;
- ruff check .;
- git diff --check.

Closes #5640.

Coding agent: Codex
Codex version: codex-cli 0.144.1
Model: gpt-5.6-sol
Reasoning effort: xhigh


<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit

* **Bug Fixes**
* Improved energy reduction accuracy by using backend-specific precision
for accumulation and per-frame/per-node division, instead of a fixed
global dtype.
* Ensured intensive/extensive reductions consistently cast reducible
values—and any provided masks or computed divisors—to the correct
backend dtype before summing/averaging.

* **Tests**
* Added Torch backend-preservation tests for both dense and graph
reduction paths, covering intensive vs extensive behavior with and
without masks and validating dtype promotion and correctness.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Co-authored-by: njzjz-bot <njzjz-bot@users.noreply.github.com>
Co-authored-by: njzjz-bot <njzjz.bot@gmail.com>

Author: njzjz

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