From 764d4ca6963b707f709f7fe8d0a83d9f24511dba Mon Sep 17 00:00:00 2001 From: njzjz-bot Date: Wed, 15 Jul 2026 00:05:27 +0800 Subject: [PATCH 1/3] fix(dpmodel): keep output reductions on active backend Resolve the configured energy precision through the active array namespace before casting dense or graph fitting outputs. Use namespace casts and reductions throughout so Torch, JAX, and other array-API inputs do not call NumPy-only methods or leave their device. Add focused Torch regressions for dense extensive, dense masked-intensive, and ragged graph reductions. Fixes #5640. Coding-Agent: Codex Codex-Version: codex-cli 0.144.1 Model: gpt-5.6-sol Reasoning-Effort: xhigh --- deepmd/dpmodel/model/edge_transform_output.py | 15 ++-- deepmd/dpmodel/model/transform_output.py | 19 ++-- .../common/dpmodel/test_transform_output.py | 89 +++++++++++++++++++ 3 files changed, 111 insertions(+), 12 deletions(-) create mode 100644 source/tests/common/dpmodel/test_transform_output.py diff --git a/deepmd/dpmodel/model/edge_transform_output.py b/deepmd/dpmodel/model/edge_transform_output.py index f9cc49f874..3a25649e3a 100644 --- a/deepmd/dpmodel/model/edge_transform_output.py +++ b/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, @@ -73,6 +75,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) model_ret = dict(fit_ret.items()) @@ -81,15 +88,13 @@ 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) redu = segment_sum(vv_e, frame_id, nf) # (nf, *shape) if vdef.intensive: if mask is not None: - cnt = segment_sum( - xp.astype(mask, GLOBAL_ENER_FLOAT_PRECISION), frame_id, nf - ) + cnt = segment_sum(xp.astype(mask, energy_dtype), 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: diff --git a/deepmd/dpmodel/model/transform_output.py b/deepmd/dpmodel/model/transform_output.py index b8ca99469f..2d358ba7c5 100644 --- a/deepmd/dpmodel/model/transform_output.py +++ b/deepmd/dpmodel/model/transform_output.py @@ -2,7 +2,6 @@ import array_api_compat -import numpy as np from deepmd.dpmodel.array_api import ( Array, @@ -10,6 +9,8 @@ ) 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 diff --git a/source/tests/common/dpmodel/test_transform_output.py b/source/tests/common/dpmodel/test_transform_output.py new file mode 100644 index 0000000000..ff4be15e7b --- /dev/null +++ b/source/tests/common/dpmodel/test_transform_output.py @@ -0,0 +1,89 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Backend-preservation tests for fitting-output reductions.""" + +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, +) + + +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``.""" + import torch + + 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.""" + import torch + + 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_graph_torch_mask_count_uses_backend_dtype() -> None: + """Graph reductions must translate the NumPy precision to ``torch.dtype``.""" + import torch + + 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) + ) From bb988aea76b7a913111434b3ccfcf54265124cb3 Mon Sep 17 00:00:00 2001 From: njzjz-bot Date: Wed, 15 Jul 2026 01:30:28 +0800 Subject: [PATCH 2/3] test(dpmodel): skip transform tests without torch Collect the backend-preservation regressions only when the optional Torch dependency is available. CI installs Torch, while lightweight source-test environments can now skip the file cleanly instead of failing during test execution. Coding-Agent: Codex Codex-Version: codex-cli 0.144.1 Model: gpt-5.6-sol Reasoning-Effort: xhigh --- source/tests/common/dpmodel/test_transform_output.py | 10 ++++------ 1 file changed, 4 insertions(+), 6 deletions(-) diff --git a/source/tests/common/dpmodel/test_transform_output.py b/source/tests/common/dpmodel/test_transform_output.py index ff4be15e7b..96aee6cda2 100644 --- a/source/tests/common/dpmodel/test_transform_output.py +++ b/source/tests/common/dpmodel/test_transform_output.py @@ -1,6 +1,8 @@ # 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, ) @@ -15,6 +17,8 @@ NeighborGraph, ) +torch = pytest.importorskip("torch") + def _output_def(*, intensive: bool) -> FittingOutputDef: """Build the smallest reducible fitting definition used by these tests.""" @@ -34,8 +38,6 @@ def _output_def(*, intensive: bool) -> FittingOutputDef: def test_dense_torch_extensive_reduction_stays_on_backend() -> None: """Dense reduction must not rely on NumPy-style ``Tensor.astype``.""" - import torch - atomic = torch.tensor([[[1.0], [2.0], [3.0]]], dtype=torch.float32) coord = torch.zeros((1, 3, 3), dtype=torch.float32) @@ -50,8 +52,6 @@ def test_dense_torch_extensive_reduction_stays_on_backend() -> None: def test_dense_torch_mask_count_uses_energy_dtype() -> None: """The intensive divisor must be reduced by Torch in energy precision.""" - import torch - 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]]) @@ -67,8 +67,6 @@ def test_dense_torch_mask_count_uses_energy_dtype() -> None: def test_graph_torch_mask_count_uses_backend_dtype() -> None: """Graph reductions must translate the NumPy precision to ``torch.dtype``.""" - import torch - graph = NeighborGraph( n_node=torch.tensor([2, 1], dtype=torch.int64), edge_index=torch.empty((2, 0), dtype=torch.int64), From d3ed641fa7a92ff83acc5eb72c1490dd65ab908f Mon Sep 17 00:00:00 2001 From: njzjz-bot Date: Thu, 16 Jul 2026 17:38:05 +0800 Subject: [PATCH 3/3] test(dpmodel): cover backend output reduction branches Exercise dense and graph intensive reductions without masks and the graph extensive path on Torch, covering the backend-native dtype branches requested in review. Coding-Agent: Codex Codex-Version: codex-cli 0.144.4 Model: gpt-5.6-sol Reasoning-Effort: xhigh --- .../common/dpmodel/test_transform_output.py | 56 +++++++++++++++++++ 1 file changed, 56 insertions(+) diff --git a/source/tests/common/dpmodel/test_transform_output.py b/source/tests/common/dpmodel/test_transform_output.py index 96aee6cda2..a406d4e6ff 100644 --- a/source/tests/common/dpmodel/test_transform_output.py +++ b/source/tests/common/dpmodel/test_transform_output.py @@ -65,6 +65,20 @@ def test_dense_torch_mask_count_uses_energy_dtype() -> None: 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( @@ -85,3 +99,45 @@ def test_graph_torch_mask_count_uses_backend_dtype() -> None: 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) + )