diff --git a/deepmd/dpmodel/model/edge_transform_output.py b/deepmd/dpmodel/model/edge_transform_output.py index ce50eba87b..a9ff3a0f1f 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, @@ -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: 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..a406d4e6ff --- /dev/null +++ b/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) + )