From 1e74dddee07ace7336ba590daa9910f1e643719f Mon Sep 17 00:00:00 2001 From: njzjz-bot Date: Fri, 17 Jul 2026 07:46:43 +0800 Subject: [PATCH 1/7] fix(dpmodel): remap virtual type embeddings Map negative virtual atom types to the explicit padding embedding row before array-api gathers and type-pair indexing. Apply the shared boundary across DPA1 through DPA4, SeZM, and SeTTebd descriptor paths. Cover NumPy and array_api_strict embedding gathers, DPA1 concat and strip modes, SeZM, and an independent SeTTebd strip pair-index regression comparing -1 with the explicit padding type. Coding-Agent: Codex Codex-Version: codex-cli 0.144.4 Model: gpt-5.6-sol Reasoning-Effort: xhigh --- deepmd/dpmodel/descriptor/dpa1.py | 17 ++- deepmd/dpmodel/descriptor/dpa2.py | 3 +- deepmd/dpmodel/descriptor/dpa3.py | 6 +- .../dpmodel/descriptor/dpa4_nn/embedding.py | 2 + deepmd/dpmodel/descriptor/se_t_tebd.py | 5 +- deepmd/dpmodel/utils/type_embed.py | 25 ++++ .../dpmodel/test_type_embedding_virtual.py | 135 ++++++++++++++++++ 7 files changed, 183 insertions(+), 10 deletions(-) create mode 100644 source/tests/common/dpmodel/test_type_embedding_virtual.py diff --git a/deepmd/dpmodel/descriptor/dpa1.py b/deepmd/dpmodel/descriptor/dpa1.py index 3659e30776..13df86fc8d 100644 --- a/deepmd/dpmodel/descriptor/dpa1.py +++ b/deepmd/dpmodel/descriptor/dpa1.py @@ -52,6 +52,8 @@ ) from deepmd.dpmodel.utils.type_embed import ( TypeEmbedNet, + remap_atype_to_padding, + take_type_embedding, ) from deepmd.dpmodel.utils.update_sel import ( UpdateSel, @@ -755,7 +757,7 @@ def _call_dense( type_embedding = self.type_embedding.call() # nf x nall x tebd_dim atype_embd_ext = xp.reshape( - xp.take(type_embedding, xp.reshape(atype_ext, (-1,)), axis=0), + take_type_embedding(type_embedding, xp.reshape(atype_ext, (-1,))), (nf, nall, self.tebd_dim), ) # nfnl x tebd_dim @@ -850,7 +852,7 @@ def call_graph( # gradient so the tebd net never trains; type_embedding already lives # on the model device, so the device cast was redundant anyway. atype_local = xp.asarray(atype, device=dev) - atype_embd = xp.take(type_embedding, atype_local, axis=0) # (N, tebd_dim) + atype_embd = take_type_embedding(type_embedding, atype_local) grrg = xp.concat([grrg, atype_embd], axis=-1) if in_dtype != prec: grrg = xp.astype(grrg, in_dtype) @@ -1646,6 +1648,7 @@ def call( # Gather neighbor types: (nf, nall) -> (nf, nloc*nnei) nei_type = xp_take_along_axis(atype_ext, nlist_2d, axis=1) nei_type = xp.reshape(nei_type, (-1,)) # (nf * nloc * nnei,) + nei_type = remap_atype_to_padding(nei_type, ntypes_with_padding) # (nf x nl x nnei) x ng nei_type_index = xp.tile(xp.reshape(nei_type, (-1, 1)), (1, ng)) if self.type_one_side: @@ -1660,9 +1663,11 @@ def call( # (nf x nl x nnei) x ng gg_t = xp_take_along_axis(tt_full, nei_type_index, axis=0) else: + center_type = remap_atype_to_padding(atype, ntypes_with_padding) idx_i = xp.reshape( xp.tile( - (xp.reshape(atype, (-1, 1)) * ntypes_with_padding), (1, nnei) + (xp.reshape(center_type, (-1, 1)) * ntypes_with_padding), + (1, nnei), ), (-1,), ) @@ -1850,9 +1855,9 @@ def call_graph( # under torch and severs the type-embedding weight gradient (the tebd # net would never train); type_embedding already lives on the device. tebd = type_embedding - atype_embd_nlist = xp.take(tebd, nei_type, axis=0) # (E, tebd_dim) + atype_embd_nlist = take_type_embedding(tebd, nei_type) if not self.type_one_side: - atype_embd_nnei = xp.take(tebd, center_type, axis=0) # (E, tebd_dim) + atype_embd_nnei = take_type_embedding(tebd, center_type) ss = xp.concat([ss, atype_embd_nlist, atype_embd_nnei], axis=-1) else: ss = xp.concat([ss, atype_embd_nlist], axis=-1) @@ -1929,6 +1934,8 @@ def _graph_edge_gg_strip( xp = array_api_compat.array_namespace(ss) nt = self.tebd_dim ntypes_with_padding = type_embedding.shape[0] + center_type = remap_atype_to_padding(center_type, ntypes_with_padding) + nei_type = remap_atype_to_padding(nei_type, ntypes_with_padding) # geometric net on the radial channel only (dense: gg_s = cal_g(ss_scalar)) gg_s = self.embeddings[0].call(ss) # (E, ng) if self.type_one_side: diff --git a/deepmd/dpmodel/descriptor/dpa2.py b/deepmd/dpmodel/descriptor/dpa2.py index 08f68849d2..f5968b6883 100644 --- a/deepmd/dpmodel/descriptor/dpa2.py +++ b/deepmd/dpmodel/descriptor/dpa2.py @@ -38,6 +38,7 @@ ) from deepmd.dpmodel.utils.type_embed import ( TypeEmbedNet, + take_type_embedding, ) from deepmd.dpmodel.utils.update_sel import ( UpdateSel, @@ -893,7 +894,7 @@ def call( type_embedding = self.type_embedding.call() # repinit g1_ext = xp.reshape( - xp.take(type_embedding, xp.reshape(atype_ext, (-1,)), axis=0), + take_type_embedding(type_embedding, xp.reshape(atype_ext, (-1,))), (nframes, nall, self.tebd_dim), ) g1_inp = xp_take_first_n(g1_ext, 1, nloc) diff --git a/deepmd/dpmodel/descriptor/dpa3.py b/deepmd/dpmodel/descriptor/dpa3.py index 76124fa544..0dfe671438 100644 --- a/deepmd/dpmodel/descriptor/dpa3.py +++ b/deepmd/dpmodel/descriptor/dpa3.py @@ -28,6 +28,7 @@ ) from deepmd.dpmodel.utils.type_embed import ( TypeEmbedNet, + take_type_embedding, ) from deepmd.dpmodel.utils.update_sel import ( UpdateSel, @@ -707,16 +708,15 @@ def call( type_embedding = self.type_embedding.call() if self.use_loc_mapping: node_ebd_ext = xp.reshape( - xp.take( + take_type_embedding( type_embedding, xp.reshape(xp_take_first_n(atype_ext, 1, nloc), (-1,)), - axis=0, ), (nframes, nloc, self.tebd_dim), ) else: node_ebd_ext = xp.reshape( - xp.take(type_embedding, xp.reshape(atype_ext, (-1,)), axis=0), + take_type_embedding(type_embedding, xp.reshape(atype_ext, (-1,))), (nframes, nall, self.tebd_dim), ) diff --git a/deepmd/dpmodel/descriptor/dpa4_nn/embedding.py b/deepmd/dpmodel/descriptor/dpa4_nn/embedding.py index 238f904a8c..6f1f02c64a 100644 --- a/deepmd/dpmodel/descriptor/dpa4_nn/embedding.py +++ b/deepmd/dpmodel/descriptor/dpa4_nn/embedding.py @@ -146,6 +146,8 @@ def call(self, atype: Any) -> Any: # torch.embedding gather: flatten the indices to int64, take the rows, # then restore the original index shape. index = xp.astype(xp.reshape(atype, (-1,)), xp.int64) + if self.padding: + index = xp.where(index >= 0, index, xp.full_like(index, self.ntypes)) out = xp.take(weight, index, axis=0) return xp.reshape(out, (*atype.shape, self.embed_dim)) diff --git a/deepmd/dpmodel/descriptor/se_t_tebd.py b/deepmd/dpmodel/descriptor/se_t_tebd.py index cb174896cb..6c5ab4c04d 100644 --- a/deepmd/dpmodel/descriptor/se_t_tebd.py +++ b/deepmd/dpmodel/descriptor/se_t_tebd.py @@ -36,6 +36,8 @@ ) from deepmd.dpmodel.utils.type_embed import ( TypeEmbedNet, + remap_atype_to_padding, + take_type_embedding, ) from deepmd.dpmodel.utils.update_sel import ( UpdateSel, @@ -398,7 +400,7 @@ def call( type_embedding = self.type_embedding.call() # nf x nall x tebd_dim atype_embd_ext = xp.reshape( - xp.take(type_embedding, xp.reshape(atype_ext, (-1,)), axis=0), + take_type_embedding(type_embedding, xp.reshape(atype_ext, (-1,))), (nf, nall, self.tebd_dim), ) # nfnl x tebd_dim @@ -933,6 +935,7 @@ def call( nei_type = xp_take_along_axis(atype_ext, nlist_index, axis=1) # nfnl x nnei nei_type = xp.reshape(nei_type, (nf * nloc, nnei)) + nei_type = remap_atype_to_padding(nei_type, ntypes_with_padding) # nfnl x nnei x nnei nei_type_i = xp.tile(nei_type[:, :, np.newaxis], (1, 1, nnei)) diff --git a/deepmd/dpmodel/utils/type_embed.py b/deepmd/dpmodel/utils/type_embed.py index f945d0f801..f151c98faf 100644 --- a/deepmd/dpmodel/utils/type_embed.py +++ b/deepmd/dpmodel/utils/type_embed.py @@ -31,6 +31,31 @@ def _array_device_or_none(array: Array) -> Any: return None +def remap_atype_to_padding(atype: Array, ntypes_with_padding: int) -> Array: + """Map negative placeholder types to a table's final padding row.""" + xp = array_api_compat.array_namespace(atype) + return xp.where( + atype >= 0, + atype, + xp.full_like(atype, ntypes_with_padding - 1), + ) + + +def take_type_embedding(type_embedding: Array, atype: Array) -> Array: + """Gather type embeddings, mapping virtual atom types to the padding row. + + Descriptor type-embedding tables append an all-zero final row for virtual + atoms. Negative placeholder types must select that row explicitly because + negative gather indices either wrap or fail depending on the array backend. + """ + # The caller's atom-type array determines the active backend. Model + # conversion keeps the embedding table in that same namespace while + # preserving trainable tensors and their gradients. + xp = array_api_compat.array_namespace(atype) + safe_atype = remap_atype_to_padding(atype, type_embedding.shape[0]) + return xp.take(type_embedding, xp.astype(safe_atype, xp.int64), axis=0) + + class TypeEmbedNet(NativeOP): r"""Type embedding network. diff --git a/source/tests/common/dpmodel/test_type_embedding_virtual.py b/source/tests/common/dpmodel/test_type_embedding_virtual.py new file mode 100644 index 0000000000..7fe9aa1b82 --- /dev/null +++ b/source/tests/common/dpmodel/test_type_embedding_virtual.py @@ -0,0 +1,135 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Regression tests for virtual types in dpmodel embedding gathers.""" + +import array_api_strict +import numpy as np +import pytest + +from deepmd.dpmodel.common import ( + to_numpy_array, +) +from deepmd.dpmodel.descriptor.dpa1 import ( + DescrptDPA1, +) +from deepmd.dpmodel.descriptor.dpa4_nn.embedding import ( + SeZMTypeEmbedding, +) +from deepmd.dpmodel.descriptor.se_t_tebd import ( + DescrptSeTTebd, +) +from deepmd.dpmodel.utils.type_embed import ( + take_type_embedding, +) +from source.tests.array_api_strict.common import ( + convert_array_api_strict_value, +) + + +@pytest.mark.parametrize("namespace", [np, array_api_strict]) +def test_padded_type_embedding_maps_virtual_type(namespace) -> None: + """Negative types select the explicit final zero row on every backend.""" + table = namespace.asarray( + np.array([[1.0, 2.0], [3.0, 4.0], [0.0, 0.0]], dtype=np.float64) + ) + atype = namespace.asarray(np.array([0, -1, 1], dtype=np.int64)) + + actual = take_type_embedding(table, atype) + + np.testing.assert_array_equal( + to_numpy_array(actual), + [[1.0, 2.0], [0.0, 0.0], [3.0, 4.0]], + ) + + +@pytest.mark.parametrize("namespace", [np, array_api_strict]) +def test_sezm_padded_embedding_maps_virtual_type(namespace) -> None: + """SeZM uses the same padding-row contract at its gather boundary.""" + embedding = SeZMTypeEmbedding(ntypes=2, embed_dim=2, padding=True, seed=1) + embedding.adam_type_embedding = np.array( + [[1.0, 2.0], [3.0, 4.0], [0.0, 0.0]], dtype=np.float64 + ) + atype = namespace.asarray(np.array([0, -1, 1], dtype=np.int64)) + + actual = embedding(atype) + + np.testing.assert_array_equal( + to_numpy_array(actual), + [[1.0, 2.0], [0.0, 0.0], [3.0, 4.0]], + ) + + +@pytest.mark.parametrize( + ("tebd_input_mode", "type_one_side"), + [("concat", True), ("strip", True), ("strip", False)], +) +def test_dpa1_strict_virtual_type_matches_explicit_padding( + tebd_input_mode: str, type_one_side: bool +) -> None: + """Direct descriptor calls remap virtual types before all gather modes.""" + descriptor = convert_array_api_strict_value( + DescrptDPA1( + rcut=4.0, + rcut_smth=0.5, + sel=[2, 2], + ntypes=2, + attn_layer=0, + axis_neuron=2, + neuron=[6, 12], + tebd_input_mode=tebd_input_mode, + type_one_side=type_one_side, + ) + ) + coord = array_api_strict.asarray( + np.array([[[0.0, 0.0, 0.0], [1.0, 0.0, 0.0], [9.0, 9.0, 9.0]]]) + ) + nlist = array_api_strict.asarray( + np.array([[[1, -1, -1, -1], [0, -1, -1, -1]]], dtype=np.int64) + ) + virtual_atype = array_api_strict.asarray(np.array([[0, 1, -1]], dtype=np.int64)) + padding_atype = array_api_strict.asarray(np.array([[0, 1, 2]], dtype=np.int64)) + + actual = descriptor._call_dense(coord, virtual_atype, nlist) + expected = descriptor._call_dense(coord, padding_atype, nlist) + + for actual_value, expected_value in zip(actual, expected, strict=True): + if actual_value is not None: + np.testing.assert_allclose( + to_numpy_array(actual_value), to_numpy_array(expected_value) + ) + + +def test_se_t_tebd_strip_strict_virtual_type_matches_explicit_padding() -> None: + """Strip-mode pair indices remap virtual neighbors to the padding type.""" + descriptor = convert_array_api_strict_value( + DescrptSeTTebd( + rcut=4.0, + rcut_smth=0.5, + sel=2, + ntypes=2, + neuron=[4, 8], + tebd_dim=2, + tebd_input_mode="strip", + concat_output_tebd=False, + seed=7, + ) + ) + coord = array_api_strict.asarray( + np.array( + [[[0.0, 0.0, 0.0], [1.0, 0.0, 0.0], [0.0, 1.0, 0.0]]], + dtype=np.float64, + ) + ) + # Atom 2 participates in both local environments so the strip-mode + # type-pair lookup must consume its remapped padding index. + nlist = array_api_strict.asarray(np.array([[[1, 2], [0, 2]]], dtype=np.int64)) + virtual_atype = array_api_strict.asarray(np.array([[0, 1, -1]], dtype=np.int64)) + padding_atype = array_api_strict.asarray(np.array([[0, 1, 2]], dtype=np.int64)) + + actual = descriptor(coord, virtual_atype, nlist) + expected = descriptor(coord, padding_atype, nlist) + + for actual_value, expected_value in zip(actual, expected, strict=True): + if actual_value is not None: + np.testing.assert_allclose( + to_numpy_array(actual_value), to_numpy_array(expected_value) + ) From b811beb6ada6a31feb27ee94b6b921aa1a8a8970 Mon Sep 17 00:00:00 2001 From: njzjz-bot Date: Sat, 1 Aug 2026 23:49:48 +0800 Subject: [PATCH 2/7] fix(descriptor): complete virtual type remapping Apply the padding-row convention to strip pair indices across dpmodel, pt_expt, PyTorch, and Paddle implementations. Keep real-type-only exclusion and normalization lookups clamped separately, document the sentinel invariant, and add focused DPA1/DPA2/DPA3 and Torch regressions. Coding-Agent: Codex Codex-Version: codex-cli 0.144.6 Model: gpt-5.6-sol Reasoning-Effort: xhigh --- deepmd/dpmodel/descriptor/dpa1.py | 9 +- .../dpmodel/descriptor/dpa4_nn/embedding.py | 5 +- deepmd/dpmodel/utils/type_embed.py | 52 +++++- deepmd/pd/model/descriptor/se_atten.py | 14 +- deepmd/pd/model/descriptor/se_t_tebd.py | 5 + deepmd/pt/model/descriptor/se_atten.py | 15 +- deepmd/pt/model/descriptor/se_t_tebd.py | 5 + deepmd/pt_expt/descriptor/dpa1.py | 21 ++- deepmd/pt_expt/descriptor/dpa2.py | 6 +- deepmd/pt_expt/descriptor/se_t_tebd.py | 7 +- .../dpmodel/test_type_embedding_virtual.py | 172 ++++++++++++++++-- .../pt/model/test_virtual_type_embedding.py | 142 +++++++++++++++ source/tests/pt_expt/descriptor/test_dpa1.py | 67 +++++++ 13 files changed, 493 insertions(+), 27 deletions(-) create mode 100644 source/tests/pt/model/test_virtual_type_embedding.py diff --git a/deepmd/dpmodel/descriptor/dpa1.py b/deepmd/dpmodel/descriptor/dpa1.py index 13df86fc8d..45c589cf3d 100644 --- a/deepmd/dpmodel/descriptor/dpa1.py +++ b/deepmd/dpmodel/descriptor/dpa1.py @@ -1822,22 +1822,27 @@ def call_graph( # value so the kernel stays jit/export-traceable (no concretize of n_node). n_total = atype.shape[0] atype = xp.asarray(atype, device=dev) + # Padded embedding tables reserve their final row, whereas exclusion + # and normalization tables contain only real types. Keep both forms so + # each downstream lookup receives the sentinel convention it expects. + safe_real_atype = xp.where(atype >= 0, atype, xp.zeros_like(atype)) # descriptor-level pair exclusion: same canonical transform as the # model-level ``pair_exclude_types`` (decision #18). Masked edges # contribute zero to every segment_sum below; the dense path's # nlist-erasure + env-mat zeroing is reproduced exactly. # apply_pair_exclusion is a no-op when self.emask has no exclusions. - graph = apply_pair_exclusion(graph, atype, self.emask) + graph = apply_pair_exclusion(graph, safe_real_atype, self.emask) src = graph.edge_index[0, :] dst = graph.edge_index[1, :] center_type = xp.take(atype, dst, axis=0) # (E,) nei_type = xp.take(atype, src, axis=0) # (E,) + center_type_for_stats = xp.take(safe_real_atype, dst, axis=0) # per-edge env-mat 4-vector, normalized by the center (dst) atom type. # self.mean/self.stddev are slot-independent (ntypes, nnei, 4); slot 0 is # the canonical per-type vector. rr, sw_e = edge_env_mat( graph.edge_vec, - center_type, + center_type_for_stats, self.mean[:, 0, :], self.stddev[:, 0, :], self.rcut, diff --git a/deepmd/dpmodel/descriptor/dpa4_nn/embedding.py b/deepmd/dpmodel/descriptor/dpa4_nn/embedding.py index 6f1f02c64a..60d2b59dbb 100644 --- a/deepmd/dpmodel/descriptor/dpa4_nn/embedding.py +++ b/deepmd/dpmodel/descriptor/dpa4_nn/embedding.py @@ -39,6 +39,9 @@ from deepmd.dpmodel.utils.network import ( NativeLayer, ) +from deepmd.dpmodel.utils.type_embed import ( + remap_atype_to_padding, +) from deepmd.dpmodel.utils.seed import ( child_seed, ) @@ -147,7 +150,7 @@ def call(self, atype: Any) -> Any: # then restore the original index shape. index = xp.astype(xp.reshape(atype, (-1,)), xp.int64) if self.padding: - index = xp.where(index >= 0, index, xp.full_like(index, self.ntypes)) + index = remap_atype_to_padding(index, self.ntypes + 1) out = xp.take(weight, index, axis=0) return xp.reshape(out, (*atype.shape, self.embed_dim)) diff --git a/deepmd/dpmodel/utils/type_embed.py b/deepmd/dpmodel/utils/type_embed.py index f151c98faf..5aaaaac6ba 100644 --- a/deepmd/dpmodel/utils/type_embed.py +++ b/deepmd/dpmodel/utils/type_embed.py @@ -32,7 +32,29 @@ def _array_device_or_none(array: Array) -> Any: def remap_atype_to_padding(atype: Array, ntypes_with_padding: int) -> Array: - """Map negative placeholder types to a table's final padding row.""" + """Map negative placeholder types to a padded table's final row. + + Parameters + ---------- + atype : Array + Atom-type indices. Negative entries denote virtual or padding atoms. + ntypes_with_padding : int + Number of rows in a table that reserves its final row for padding. + + Returns + ------- + Array + Atom-type indices with every negative entry replaced by + ``ntypes_with_padding - 1``. + + Notes + ----- + This sentinel convention is valid only for tables that explicitly include + a final padding row, such as descriptor type-embedding and type-pair + tables. It must not be used for real-type-only tables such as ``davg``, + ``dstd``, or spin masks; virtual entries must be masked or clamped to a + valid real type before indexing those tables. + """ xp = array_api_compat.array_namespace(atype) return xp.where( atype >= 0, @@ -44,9 +66,31 @@ def remap_atype_to_padding(atype: Array, ntypes_with_padding: int) -> Array: def take_type_embedding(type_embedding: Array, atype: Array) -> Array: """Gather type embeddings, mapping virtual atom types to the padding row. - Descriptor type-embedding tables append an all-zero final row for virtual - atoms. Negative placeholder types must select that row explicitly because - negative gather indices either wrap or fail depending on the array backend. + Parameters + ---------- + type_embedding : Array + Type-embedding table whose final row is reserved for virtual or + padding atoms. + atype : Array + Atom-type indices with arbitrary shape. Negative entries denote + virtual or padding atoms. + + Returns + ------- + Array + Gathered embeddings with shape ``(*atype.shape, + type_embedding.shape[-1])``. + + Notes + ----- + ``TypeEmbedNet`` reconstructs a literal zero padding row on every call. + ``SeZMTypeEmbedding`` stores its reserved row in the trainable embedding + array and initializes it to zero. This helper guarantees selection of the + reserved row; the table implementation remains responsible for keeping + that row neutral. + + Negative placeholder types must be remapped explicitly because negative + gather indices either wrap or fail depending on the array backend. """ # The caller's atom-type array determines the active backend. Model # conversion keeps the embedding table in that same namespace while diff --git a/deepmd/pd/model/descriptor/se_atten.py b/deepmd/pd/model/descriptor/se_atten.py index 33d8e8d4cf..4304b4bd42 100644 --- a/deepmd/pd/model/descriptor/se_atten.py +++ b/deepmd/pd/model/descriptor/se_atten.py @@ -582,6 +582,13 @@ def forward( nei_type = paddle.take_along_axis( extended_atype, indices=nlist_index, axis=1, broadcast=False ) + # Padded embedding tables reserve their final row for virtual + # atoms; remap explicitly before pair-index arithmetic. + nei_type = paddle.where( + nei_type >= 0, + nei_type, + paddle.full_like(nei_type, ntypes_with_padding - 1), + ) # (nf x nl x nnei) x ng nei_type_index = nei_type.reshape([-1, 1]).expand([-1, ng]).to(paddle.int64) if self.type_one_side: @@ -591,8 +598,13 @@ def forward( tt_full, indices=nei_type_index, axis=0, broadcast=False ) else: + center_type = paddle.where( + atype >= 0, + atype, + paddle.full_like(atype, ntypes_with_padding - 1), + ) idx_i = paddle.tile( - atype.reshape([-1, 1]) * ntypes_with_padding, [1, nnei] + center_type.reshape([-1, 1]) * ntypes_with_padding, [1, nnei] ).reshape([-1]) idx_j = nei_type.reshape([-1]) # (nf x nl x nnei) x ng diff --git a/deepmd/pd/model/descriptor/se_t_tebd.py b/deepmd/pd/model/descriptor/se_t_tebd.py index 0c16eb0ef0..13f78217e0 100644 --- a/deepmd/pd/model/descriptor/se_t_tebd.py +++ b/deepmd/pd/model/descriptor/se_t_tebd.py @@ -929,6 +929,11 @@ def forward( ) # nfnl x nnei nei_type = nei_type.reshape([nfnl, nnei]) + nei_type = paddle.where( + nei_type >= 0, + nei_type, + paddle.full_like(nei_type, ntypes_with_padding - 1), + ) # nfnl x nnei x nnei nei_type_i = nei_type.unsqueeze(2).expand([-1, -1, nnei]) nei_type_j = nei_type.unsqueeze(1).expand([-1, nnei, -1]) diff --git a/deepmd/pt/model/descriptor/se_atten.py b/deepmd/pt/model/descriptor/se_atten.py index 808515702a..461954281b 100644 --- a/deepmd/pt/model/descriptor/se_atten.py +++ b/deepmd/pt/model/descriptor/se_atten.py @@ -692,6 +692,14 @@ def forward( nlist_index = nlist.reshape(nb, nloc * nnei) # nf x (nl x nnei) nei_type = torch.gather(extended_atype, dim=1, index=nlist_index) + # Only padded type/type-pair tables use the final-row sentinel. + # Remap before both one- and two-side indexing so negative virtual + # types cannot wrap into an unrelated pair row. + nei_type = torch.where( + nei_type >= 0, + nei_type, + torch.full_like(nei_type, ntypes_with_padding - 1), + ) # Per-edge row index into the (padded) type-pair embedding table. if self.type_one_side: if self.tebd_compress: @@ -701,8 +709,13 @@ def forward( tt_full = self.filter_layers_strip.networks[0](type_embedding) tebd_idx = nei_type.view(-1).to(torch.long) else: + center_type = torch.where( + atype >= 0, + atype, + torch.full_like(atype, ntypes_with_padding - 1), + ) idx_i = torch.tile( - atype.reshape(-1, 1) * ntypes_with_padding, [1, nnei] + center_type.reshape(-1, 1) * ntypes_with_padding, [1, nnei] ).view(-1) tebd_idx = (idx_i + nei_type.view(-1)).to(torch.long) if self.tebd_compress: diff --git a/deepmd/pt/model/descriptor/se_t_tebd.py b/deepmd/pt/model/descriptor/se_t_tebd.py index 6937bb99e8..10302b71f6 100644 --- a/deepmd/pt/model/descriptor/se_t_tebd.py +++ b/deepmd/pt/model/descriptor/se_t_tebd.py @@ -1021,6 +1021,11 @@ def forward( nei_type = torch.gather(extended_atype, dim=1, index=nlist_index) # nfnl x nnei nei_type = nei_type.reshape(nfnl, nnei) + nei_type = torch.where( + nei_type >= 0, + nei_type, + torch.full_like(nei_type, ntypes_with_padding - 1), + ) # nfnl x nnei x nnei nei_type_i = nei_type.unsqueeze(2).expand([-1, -1, nnei]) nei_type_j = nei_type.unsqueeze(1).expand([-1, nnei, -1]) diff --git a/deepmd/pt_expt/descriptor/dpa1.py b/deepmd/pt_expt/descriptor/dpa1.py index c1520402fb..a1202e9795 100644 --- a/deepmd/pt_expt/descriptor/dpa1.py +++ b/deepmd/pt_expt/descriptor/dpa1.py @@ -13,6 +13,9 @@ from deepmd.dpmodel.utils.env_mat_stat import ( merge_env_stat, ) +from deepmd.dpmodel.utils.type_embed import ( + remap_atype_to_padding, +) from deepmd.kernels.cuda.dpa1.graph_compress import ( dpa1_graph_compress, ) @@ -167,9 +170,10 @@ def _strip_pair_index( ntypes_with_padding = type_embedding.shape[0] nlist_index = nlist_masked.view(nf, nloc * nnei) nei_type = torch.gather(atype_ext, dim=1, index=nlist_index) + nei_type = remap_atype_to_padding(nei_type, ntypes_with_padding) if se.type_one_side: return nei_type.reshape(-1).to(torch.long) - atype = atype_ext[:, :nloc] + atype = remap_atype_to_padding(atype_ext[:, :nloc], ntypes_with_padding) idx_i = torch.tile(atype.reshape(-1, 1) * ntypes_with_padding, [1, nnei]).view(-1) return (idx_i + nei_type.reshape(-1)).to(torch.long) @@ -793,10 +797,15 @@ def _call_graph_compress_reference( u = ((length - se.rcut_smth) / (se.rcut - se.rcut_smth)).clamp(0.0, 1.0) sw = u**3 * (-6 * u**2 + 15 * u - 10) + 1.0 em = torch.cat([sw / q, ev * (sw / q**2)], dim=-1) - rr = (em - se.mean[:, 0, :][center_type]) / se.stddev[:, 0, :][center_type] + center_type_for_stats = center_type.clamp_min(0) + rr = (em - se.mean[:, 0, :][center_type_for_stats]) / se.stddev[:, 0, :][ + center_type_for_stats + ] # === Step 2. Strip type-pair gate from the precomputed table === ntypes = type_embedding.shape[0] + center_type = remap_atype_to_padding(center_type, ntypes) + nei_type = remap_atype_to_padding(nei_type, ntypes) pair_idx = nei_type if se.type_one_side else center_type * ntypes + nei_type gate = self.type_embd_data[pair_idx] if se.smooth: @@ -983,6 +992,7 @@ def _call_graph_triton( dst = graph.edge_index[1, :] center_type = atype[dst] nei_type = atype[src] + center_type_for_stats = center_type.clamp_min(0) # Per-edge env-mat 4-vector, normalized by the center (dst) atom type; # mean/stddev are slot-independent, so slot 0 is the canonical vector. # The fused operator is captured opaquely under the pt_expt trace and @@ -992,7 +1002,7 @@ def _call_graph_triton( # in ``edge_vec``); the same operator emits it when ``return_sw`` is set. rr, sw_e = _edge_env_mat_triton( graph.edge_vec, - center_type, + center_type_for_stats, se.mean[:, 0, :], se.stddev[:, 0, :], se.rcut, @@ -1013,6 +1023,8 @@ def _call_graph_triton( emb_in = ss tt = _type_pair_table(self, type_embedding) ntypes_with_padding = type_embedding.shape[0] + center_type = remap_atype_to_padding(center_type, ntypes_with_padding) + nei_type = remap_atype_to_padding(nei_type, ntypes_with_padding) if se.type_one_side: gate_idx = nei_type.to(torch.long) else: @@ -1029,6 +1041,9 @@ def _call_graph_triton( # concat embedding input: radial channel plus the neighbor (and, two- # side, center) type embeddings. Ghost type == owner type, so # gathering by the local owner reproduces the dense neighbor tebd. + ntypes_with_padding = type_embedding.shape[0] + center_type = remap_atype_to_padding(center_type, ntypes_with_padding) + nei_type = remap_atype_to_padding(nei_type, ntypes_with_padding) nlist_tebd = type_embedding[nei_type] # (E, tebd_dim) if se.type_one_side: emb_in = torch.cat([ss, nlist_tebd], dim=-1) diff --git a/deepmd/pt_expt/descriptor/dpa2.py b/deepmd/pt_expt/descriptor/dpa2.py index 01a7dfde32..e523204ef9 100644 --- a/deepmd/pt_expt/descriptor/dpa2.py +++ b/deepmd/pt_expt/descriptor/dpa2.py @@ -17,6 +17,9 @@ from deepmd.dpmodel.utils.env_mat_stat import ( merge_env_stat, ) +from deepmd.dpmodel.utils.type_embed import ( + remap_atype_to_padding, +) from deepmd.pt_expt.common import ( torch_module, ) @@ -414,12 +417,13 @@ def _compressed_repinit_forward( nlist_index = nlist_masked.view(nf, nloc_r * nnei) # nf x (nloc x nnei) nei_type = torch.gather(atype_ext, dim=1, index=nlist_index) + nei_type = remap_atype_to_padding(nei_type, ntypes_with_padding) if self.repinit.type_one_side: # (nf*nl*nnei,) -> (nf*nl*nnei, ng) gg_t = self.type_embd_data[nei_type.view(-1).to(torch.long)] else: - atype = atype_ext[:, :nloc_r] + atype = remap_atype_to_padding(atype_ext[:, :nloc_r], ntypes_with_padding) idx_i = torch.tile( atype.reshape(-1, 1) * ntypes_with_padding, [1, nnei] ).view(-1) diff --git a/deepmd/pt_expt/descriptor/se_t_tebd.py b/deepmd/pt_expt/descriptor/se_t_tebd.py index 9e358d64b6..bc576f78b2 100644 --- a/deepmd/pt_expt/descriptor/se_t_tebd.py +++ b/deepmd/pt_expt/descriptor/se_t_tebd.py @@ -12,6 +12,9 @@ from deepmd.dpmodel.utils.env_mat_stat import ( merge_env_stat, ) +from deepmd.dpmodel.utils.type_embed import ( + remap_atype_to_padding, +) from deepmd.pt_expt.common import ( torch_module, ) @@ -237,7 +240,9 @@ def _call_compressed( # nf x (nloc x nnei) nei_type = torch.gather(atype_ext, dim=1, index=nlist_index) # nfnl x nnei - nei_type = nei_type.view(nfnl, nnei) + nei_type = remap_atype_to_padding( + nei_type.view(nfnl, nnei), ntypes_with_padding + ) # nfnl x nnei x nnei nei_type_i = nei_type.unsqueeze(2).expand(-1, -1, nnei) nei_type_j = nei_type.unsqueeze(1).expand(-1, nnei, -1) diff --git a/source/tests/common/dpmodel/test_type_embedding_virtual.py b/source/tests/common/dpmodel/test_type_embedding_virtual.py index 7fe9aa1b82..b809830fec 100644 --- a/source/tests/common/dpmodel/test_type_embedding_virtual.py +++ b/source/tests/common/dpmodel/test_type_embedding_virtual.py @@ -11,6 +11,15 @@ from deepmd.dpmodel.descriptor.dpa1 import ( DescrptDPA1, ) +from deepmd.dpmodel.descriptor.dpa2 import ( + DescrptDPA2, + RepformerArgs, + RepinitArgs, +) +from deepmd.dpmodel.descriptor.dpa3 import ( + DescrptDPA3, + RepFlowArgs, +) from deepmd.dpmodel.descriptor.dpa4_nn.embedding import ( SeZMTypeEmbedding, ) @@ -25,13 +34,24 @@ ) -@pytest.mark.parametrize("namespace", [np, array_api_strict]) -def test_padded_type_embedding_maps_virtual_type(namespace) -> None: - """Negative types select the explicit final zero row on every backend.""" - table = namespace.asarray( - np.array([[1.0, 2.0], [3.0, 4.0], [0.0, 0.0]], dtype=np.float64) - ) - atype = namespace.asarray(np.array([0, -1, 1], dtype=np.int64)) +@pytest.mark.parametrize("namespace_name", ["numpy", "torch"]) +def test_padded_type_embedding_maps_virtual_type(namespace_name: str) -> None: + """Negative types select the explicit final zero row, including on Torch.""" + if namespace_name == "torch": + import torch + + namespace = torch + table = namespace.asarray( + np.array([[1.0, 2.0], [3.0, 4.0], [0.0, 0.0]], dtype=np.float64), + device="cpu", + ) + atype = namespace.asarray(np.array([0, -1, 1], dtype=np.int64), device="cpu") + else: + namespace = np + table = namespace.asarray( + np.array([[1.0, 2.0], [3.0, 4.0], [0.0, 0.0]], dtype=np.float64) + ) + atype = namespace.asarray(np.array([0, -1, 1], dtype=np.int64)) actual = take_type_embedding(table, atype) @@ -41,15 +61,21 @@ def test_padded_type_embedding_maps_virtual_type(namespace) -> None: ) -@pytest.mark.parametrize("namespace", [np, array_api_strict]) -def test_sezm_padded_embedding_maps_virtual_type(namespace) -> None: +@pytest.mark.parametrize("namespace_name", ["numpy", "torch"]) +def test_sezm_padded_embedding_maps_virtual_type(namespace_name: str) -> None: """SeZM uses the same padding-row contract at its gather boundary.""" + if namespace_name == "torch": + import torch + + namespace = torch + atype = namespace.asarray(np.array([0, -1, 1], dtype=np.int64), device="cpu") + else: + namespace = np + atype = namespace.asarray(np.array([0, -1, 1], dtype=np.int64)) embedding = SeZMTypeEmbedding(ntypes=2, embed_dim=2, padding=True, seed=1) embedding.adam_type_embedding = np.array( [[1.0, 2.0], [3.0, 4.0], [0.0, 0.0]], dtype=np.float64 ) - atype = namespace.asarray(np.array([0, -1, 1], dtype=np.int64)) - actual = embedding(atype) np.testing.assert_array_equal( @@ -80,10 +106,10 @@ def test_dpa1_strict_virtual_type_matches_explicit_padding( ) ) coord = array_api_strict.asarray( - np.array([[[0.0, 0.0, 0.0], [1.0, 0.0, 0.0], [9.0, 9.0, 9.0]]]) + np.array([[[0.0, 0.0, 0.0], [1.0, 0.0, 0.0], [0.0, 1.0, 0.0]]]) ) nlist = array_api_strict.asarray( - np.array([[[1, -1, -1, -1], [0, -1, -1, -1]]], dtype=np.int64) + np.array([[[1, 2, -1, -1], [0, 2, -1, -1]]], dtype=np.int64) ) virtual_atype = array_api_strict.asarray(np.array([[0, 1, -1]], dtype=np.int64)) padding_atype = array_api_strict.asarray(np.array([[0, 1, 2]], dtype=np.int64)) @@ -133,3 +159,123 @@ def test_se_t_tebd_strip_strict_virtual_type_matches_explicit_padding() -> None: np.testing.assert_allclose( to_numpy_array(actual_value), to_numpy_array(expected_value) ) + + +def test_dpa2_virtual_neighbor_matches_explicit_padding() -> None: + """DPA2 gathers the padding embedding for a virtual extended atom.""" + descriptor = DescrptDPA2( + ntypes=2, + repinit=RepinitArgs( + rcut=4.0, + rcut_smth=0.5, + nsel=4, + neuron=[4, 8], + axis_neuron=2, + tebd_dim=2, + tebd_input_mode="strip", + ), + repformer=RepformerArgs( + rcut=3.0, + rcut_smth=0.5, + nsel=2, + nlayers=1, + g1_dim=8, + g2_dim=4, + axis_neuron=2, + attn1_hidden=8, + attn1_nhead=2, + attn2_hidden=4, + attn2_nhead=2, + ), + concat_output_tebd=False, + seed=11, + ) + coord = np.array([[[0.0, 0.0, 0.0], [1.0, 0.0, 0.0], [0.0, 1.0, 0.0]]]) + nlist = np.array([[[1, 2, -1, -1], [0, 2, -1, -1]]], dtype=np.int64) + mapping = np.array([[0, 1, 0]], dtype=np.int64) + + actual = descriptor( + coord, + np.array([[0, 1, -1]], dtype=np.int64), + nlist, + mapping, + ) + expected = descriptor( + coord, + np.array([[0, 1, 2]], dtype=np.int64), + nlist, + mapping, + ) + + for actual_value, expected_value in zip(actual, expected, strict=True): + np.testing.assert_allclose(actual_value, expected_value) + + +@pytest.mark.parametrize("use_loc_mapping", [True, False]) +def test_dpa3_virtual_type_uses_padding_for_both_mapping_paths( + use_loc_mapping: bool, +) -> None: + """DPA3 remaps before either local-only or full extended gathering.""" + descriptor = DescrptDPA3( + ntypes=2, + repflow=RepFlowArgs( + n_dim=8, + e_dim=4, + a_dim=4, + nlayers=1, + e_rcut=4.0, + e_rcut_smth=0.5, + e_sel=4, + a_rcut=4.0, + a_rcut_smth=0.5, + a_sel=3, + axis_neuron=2, + update_angle=False, + ), + use_loc_mapping=use_loc_mapping, + seed=17, + ) + table = np.vstack( + ( + np.arange(descriptor.tebd_dim, dtype=np.float64), + np.arange(descriptor.tebd_dim, dtype=np.float64) + 10.0, + np.zeros(descriptor.tebd_dim), + ) + ) + descriptor.type_embedding.call = lambda: table + + class CaptureRepflows: + """Capture DPA3's initial node embeddings without later env lookups.""" + + node_ebd_ext: np.ndarray + + def __call__( + self, + nlist, + coord_ext, + atype_ext, + node_ebd_ext, + mapping, + comm_dict=None, + ): + self.node_ebd_ext = node_ebd_ext + nframes, nloc, nnei = nlist.shape + return ( + node_ebd_ext[:, :nloc, :], + np.zeros((nframes, nloc, nnei, 4)), + np.zeros((nframes, nloc, nnei, 3)), + np.zeros((nframes, nloc, 4, 3)), + np.zeros((nframes, nloc, nnei)), + ) + + capture = CaptureRepflows() + descriptor.repflows = capture + coord = np.array([[[0.0, 0.0, 0.0], [1.0, 0.0, 0.0], [0.0, 1.0, 0.0]]]) + nlist = np.array([[[1, 2, -1, -1], [0, 2, -1, -1]]], dtype=np.int64) + mapping = np.array([[0, 1, 0]], dtype=np.int64) + atype = np.array([[0, -1, 1] if use_loc_mapping else [0, 1, -1]], dtype=np.int64) + + descriptor(coord, atype, nlist, mapping) + + expected_rows = [0, 2] if use_loc_mapping else [0, 1, 2] + np.testing.assert_array_equal(capture.node_ebd_ext, table[expected_rows][None, ...]) diff --git a/source/tests/pt/model/test_virtual_type_embedding.py b/source/tests/pt/model/test_virtual_type_embedding.py new file mode 100644 index 0000000000..70b6eee6b3 --- /dev/null +++ b/source/tests/pt/model/test_virtual_type_embedding.py @@ -0,0 +1,142 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Regression tests for virtual types in legacy PyTorch strip descriptors.""" + +import torch + +from deepmd.pt.model.descriptor.se_atten import ( + DescrptBlockSeAtten, +) +from deepmd.pt.model.descriptor.se_t_tebd import ( + DescrptBlockSeTTebd, +) +from deepmd.pt.utils import ( + env, +) + + +def _inputs(nnei: int): + device = env.DEVICE + coord = torch.tensor( + [[[0.0, 0.0, 0.0], [1.0, 0.0, 0.0], [0.0, 1.0, 0.0]]], + dtype=torch.float64, + device=device, + ) + full_nlist = torch.tensor( + [[[1, 2, -1, -1], [0, 2, -1, -1]]], + dtype=torch.long, + device=device, + ) + virtual_atype = torch.tensor([[0, 1, -1]], dtype=torch.long, device=device) + padding_atype = torch.tensor([[0, 1, 2]], dtype=torch.long, device=device) + type_embedding = torch.tensor( + [[1.0, 2.0], [3.0, 4.0], [0.0, 0.0]], + dtype=torch.float64, + device=device, + ) + virtual_embedding = type_embedding[ + torch.where( + virtual_atype >= 0, + virtual_atype, + torch.full_like(virtual_atype, 2), + ) + ] + return ( + coord, + full_nlist[:, :, :nnei], + virtual_atype, + padding_atype, + type_embedding, + virtual_embedding, + type_embedding[padding_atype], + ) + + +def _assert_outputs_close(actual, expected) -> None: + for actual_value, expected_value in zip(actual, expected, strict=True): + if actual_value is not None: + torch.testing.assert_close(actual_value, expected_value) + + +def test_se_atten_strip_virtual_neighbor_matches_padding() -> None: + """Two-side DPA1 pair indices remap a virtual neighbor before folding.""" + ( + coord, + nlist, + virtual_atype, + padding_atype, + type_embedding, + virtual_embedding, + padding_embedding, + ) = _inputs(4) + descriptor = DescrptBlockSeAtten( + rcut=4.0, + rcut_smth=0.5, + sel=[2, 2], + ntypes=2, + attn_layer=0, + axis_neuron=2, + neuron=[6, 12], + tebd_dim=2, + tebd_input_mode="strip", + type_one_side=False, + precision="float64", + seed=1, + ).to(env.DEVICE) + + actual = descriptor( + nlist, + coord, + virtual_atype, + virtual_embedding, + type_embedding=type_embedding, + ) + expected = descriptor( + nlist, + coord, + padding_atype, + padding_embedding, + type_embedding=type_embedding, + ) + + _assert_outputs_close(actual, expected) + + +def test_se_t_tebd_strip_virtual_neighbor_matches_padding() -> None: + """SE_T type-pair indices remap virtual neighbors on both pair axes.""" + ( + coord, + nlist, + virtual_atype, + padding_atype, + type_embedding, + virtual_embedding, + padding_embedding, + ) = _inputs(2) + descriptor = DescrptBlockSeTTebd( + rcut=4.0, + rcut_smth=0.5, + sel=2, + ntypes=2, + neuron=[4, 8], + tebd_dim=2, + tebd_input_mode="strip", + precision="float64", + seed=1, + ).to(env.DEVICE) + + actual = descriptor( + nlist, + coord, + virtual_atype, + virtual_embedding, + type_embedding=type_embedding, + ) + expected = descriptor( + nlist, + coord, + padding_atype, + padding_embedding, + type_embedding=type_embedding, + ) + + _assert_outputs_close(actual, expected) diff --git a/source/tests/pt_expt/descriptor/test_dpa1.py b/source/tests/pt_expt/descriptor/test_dpa1.py index 8fd4e77598..dbd98f750d 100644 --- a/source/tests/pt_expt/descriptor/test_dpa1.py +++ b/source/tests/pt_expt/descriptor/test_dpa1.py @@ -3,6 +3,9 @@ import numpy as np import pytest import torch +from types import ( + SimpleNamespace, +) from torch.fx.experimental.proxy_tensor import ( make_fx, ) @@ -10,6 +13,7 @@ from deepmd.dpmodel.descriptor.dpa1 import DescrptDPA1 as DPDescrptDPA1 from deepmd.pt_expt.descriptor.dpa1 import ( DescrptDPA1, + _strip_pair_index, ) from deepmd.pt_expt.utils import ( env, @@ -31,6 +35,69 @@ ) +def test_strip_virtual_neighbor_matches_explicit_padding() -> None: + """Torch strip-mode pair arithmetic uses the reserved padding type.""" + device = env.DEVICE + descriptor = DescrptDPA1( + rcut=4.0, + rcut_smth=0.5, + sel=[2, 2], + ntypes=2, + attn_layer=0, + axis_neuron=2, + neuron=[6, 12], + tebd_dim=2, + tebd_input_mode="strip", + type_one_side=False, + seed=GLOBAL_SEED, + ).to(device) + coord = torch.tensor( + [[[0.0, 0.0, 0.0], [1.0, 0.0, 0.0], [0.0, 1.0, 0.0]]], + dtype=torch.float64, + device=device, + ) + nlist = torch.tensor( + [[[1, 2, -1, -1], [0, 2, -1, -1]]], + dtype=torch.long, + device=device, + ) + + actual = descriptor( + coord, + torch.tensor([[0, 1, -1]], dtype=torch.long, device=device), + nlist, + ) + expected = descriptor( + coord, + torch.tensor([[0, 1, 2]], dtype=torch.long, device=device), + nlist, + ) + + for actual_value, expected_value in zip(actual, expected, strict=True): + if actual_value is not None: + torch.testing.assert_close(actual_value, expected_value) + + +def test_strip_pair_index_remaps_virtual_types() -> None: + """The pt_expt override folds only explicit padded type indices.""" + device = env.DEVICE + descriptor = SimpleNamespace(se_atten=SimpleNamespace(type_one_side=False)) + actual = _strip_pair_index( + descriptor, + torch.tensor([[0, 1, -1]], dtype=torch.long, device=device), + torch.tensor([[1, 2], [0, 2]], dtype=torch.long, device=device), + torch.zeros((3, 2), dtype=torch.float64, device=device), + nf=1, + nloc=2, + nnei=2, + ) + + torch.testing.assert_close( + actual, + torch.tensor([1, 2, 3, 5], dtype=torch.long, device=device), + ) + + class TestDescrptDPA1(TestCaseSingleFrameWithNlist): def setup_method(self) -> None: TestCaseSingleFrameWithNlist.setUp(self) From aac6298e25feeb6208d0b5b348e47d14ae08b753 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sat, 1 Aug 2026 15:50:43 +0000 Subject: [PATCH 3/7] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- deepmd/dpmodel/descriptor/dpa4_nn/embedding.py | 6 +++--- source/tests/pt_expt/descriptor/test_dpa1.py | 7 ++++--- 2 files changed, 7 insertions(+), 6 deletions(-) diff --git a/deepmd/dpmodel/descriptor/dpa4_nn/embedding.py b/deepmd/dpmodel/descriptor/dpa4_nn/embedding.py index 60d2b59dbb..e27af3b8a2 100644 --- a/deepmd/dpmodel/descriptor/dpa4_nn/embedding.py +++ b/deepmd/dpmodel/descriptor/dpa4_nn/embedding.py @@ -39,12 +39,12 @@ from deepmd.dpmodel.utils.network import ( NativeLayer, ) -from deepmd.dpmodel.utils.type_embed import ( - remap_atype_to_padding, -) from deepmd.dpmodel.utils.seed import ( child_seed, ) +from deepmd.dpmodel.utils.type_embed import ( + remap_atype_to_padding, +) from deepmd.utils.version import ( check_version_compatibility, ) diff --git a/source/tests/pt_expt/descriptor/test_dpa1.py b/source/tests/pt_expt/descriptor/test_dpa1.py index dbd98f750d..a9ec8cde71 100644 --- a/source/tests/pt_expt/descriptor/test_dpa1.py +++ b/source/tests/pt_expt/descriptor/test_dpa1.py @@ -1,11 +1,12 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -import numpy as np -import pytest -import torch from types import ( SimpleNamespace, ) + +import numpy as np +import pytest +import torch from torch.fx.experimental.proxy_tensor import ( make_fx, ) From 4d10c3c480fa94d992933072068f8974aef5e959 Mon Sep 17 00:00:00 2001 From: njzjz-bot Date: Sun, 2 Aug 2026 00:18:41 +0800 Subject: [PATCH 4/7] fix(descriptor): sanitize virtual center statistics Clamp virtual center types before real-type-only environment-statistics lookups in PyTorch, Paddle, and pt_expt paths while retaining padding-row remapping for embedding tables. Add dense, fused-prologue, and graph regressions. Coding-Agent: Codex Codex-Version: codex-cli 0.144.6 Model: gpt-5.6-sol Reasoning-Effort: xhigh --- deepmd/pd/model/descriptor/se_atten.py | 3 +- deepmd/pd/model/descriptor/se_t_tebd.py | 3 +- deepmd/pt/model/descriptor/se_atten.py | 3 +- deepmd/pt/model/descriptor/se_t_tebd.py | 3 +- deepmd/pt_expt/descriptor/dpa1.py | 5 +- deepmd/pt_expt/descriptor/dpa2.py | 3 +- deepmd/pt_expt/descriptor/se_t_tebd.py | 3 +- .../dpmodel/test_dpa1_call_graph_block.py | 26 ++++ .../pt/model/test_virtual_type_embedding.py | 124 ++++++++++++++++++ source/tests/pt_expt/descriptor/test_dpa1.py | 46 +++++++ 10 files changed, 211 insertions(+), 8 deletions(-) diff --git a/deepmd/pd/model/descriptor/se_atten.py b/deepmd/pd/model/descriptor/se_atten.py index 4304b4bd42..411a686762 100644 --- a/deepmd/pd/model/descriptor/se_atten.py +++ b/deepmd/pd/model/descriptor/se_atten.py @@ -502,12 +502,13 @@ def forward( assert extended_atype_embd is not None nframes, nloc, nnei = nlist.shape atype = extended_atype[:, :nloc] + atype_for_env = paddle.where(atype >= 0, atype, paddle.zeros_like(atype)) nb = nframes nall = extended_coord.reshape([nb, -1, 3]).shape[1] dmatrix, diff, sw = prod_env_mat( extended_coord, nlist, - atype, + atype_for_env, self.mean, self.stddev, self.rcut, diff --git a/deepmd/pd/model/descriptor/se_t_tebd.py b/deepmd/pd/model/descriptor/se_t_tebd.py index 13f78217e0..d61294650a 100644 --- a/deepmd/pd/model/descriptor/se_t_tebd.py +++ b/deepmd/pd/model/descriptor/se_t_tebd.py @@ -850,12 +850,13 @@ def forward( assert extended_atype_embd is not None nframes, nloc, nnei = nlist.shape atype = extended_atype[:, :nloc] + atype_for_env = paddle.where(atype >= 0, atype, paddle.zeros_like(atype)) nb = nframes nall = extended_coord.reshape([nb, -1, 3]).shape[1] dmatrix, diff, sw = prod_env_mat( extended_coord, nlist, - atype, + atype_for_env, self.mean, self.stddev, self.rcut, diff --git a/deepmd/pt/model/descriptor/se_atten.py b/deepmd/pt/model/descriptor/se_atten.py index 461954281b..1a30ce72c9 100644 --- a/deepmd/pt/model/descriptor/se_atten.py +++ b/deepmd/pt/model/descriptor/se_atten.py @@ -588,12 +588,13 @@ def forward( assert extended_atype_embd is not None nframes, nloc, nnei = nlist.shape atype = extended_atype[:, :nloc] + atype_for_env = atype.clamp_min(0) nb = nframes nall = extended_coord.view(nb, -1, 3).shape[1] dmatrix, diff, sw = prod_env_mat( extended_coord, nlist, - atype, + atype_for_env, self.mean, self.stddev, self.rcut, diff --git a/deepmd/pt/model/descriptor/se_t_tebd.py b/deepmd/pt/model/descriptor/se_t_tebd.py index 10302b71f6..90a53b4b3f 100644 --- a/deepmd/pt/model/descriptor/se_t_tebd.py +++ b/deepmd/pt/model/descriptor/se_t_tebd.py @@ -932,12 +932,13 @@ def forward( assert extended_atype_embd is not None nframes, nloc, nnei = nlist.shape atype = extended_atype[:, :nloc] + atype_for_env = atype.clamp_min(0) nb = nframes nall = extended_coord.view(nb, -1, 3).shape[1] dmatrix, diff, sw = prod_env_mat( extended_coord, nlist, - atype, + atype_for_env, self.mean, self.stddev, self.rcut, diff --git a/deepmd/pt_expt/descriptor/dpa1.py b/deepmd/pt_expt/descriptor/dpa1.py index a1202e9795..4d18e2a4bc 100644 --- a/deepmd/pt_expt/descriptor/dpa1.py +++ b/deepmd/pt_expt/descriptor/dpa1.py @@ -107,6 +107,7 @@ def _env_mat( """ se = desc.se_atten nf, nloc, nnei = nlist.shape + atype_ext_for_env = atype_ext.clamp_min(0) if triton_infer_level() >= 1: # Fused env-matrix operator, captured opaquely under the pt_expt trace and # resolving to the Triton kernel at CUDA runtime; identical outputs to the @@ -114,7 +115,7 @@ def _env_mat( rr, _diff, sw = _env_mat_triton( coord_ext, nlist, - atype_ext[:, :nloc], + atype_ext_for_env[:, :nloc], se.mean[...], se.stddev[...], se.env_mat.rcut, @@ -125,7 +126,7 @@ def _env_mat( ) else: rr, _diff, sw = se.env_mat.call( - coord_ext, atype_ext, nlist, se.mean[...], se.stddev[...] + coord_ext, atype_ext_for_env, nlist, se.mean[...], se.stddev[...] ) nf, nloc, nnei, _ = rr.shape ng = se.neuron[-1] diff --git a/deepmd/pt_expt/descriptor/dpa2.py b/deepmd/pt_expt/descriptor/dpa2.py index e523204ef9..75b8f603cc 100644 --- a/deepmd/pt_expt/descriptor/dpa2.py +++ b/deepmd/pt_expt/descriptor/dpa2.py @@ -379,9 +379,10 @@ def _compressed_repinit_forward( Repinit output. shape: nf x nloc x (ng x axis_neuron) """ # env_mat: nf x nloc x nnei x 4 + atype_ext_for_env = atype_ext.clamp_min(0) rr, _diff, sw = self.repinit.env_mat.call( coord_ext, - atype_ext, + atype_ext_for_env, nlist, self.repinit.mean[...], self.repinit.stddev[...], diff --git a/deepmd/pt_expt/descriptor/se_t_tebd.py b/deepmd/pt_expt/descriptor/se_t_tebd.py index bc576f78b2..512e0ebe7a 100644 --- a/deepmd/pt_expt/descriptor/se_t_tebd.py +++ b/deepmd/pt_expt/descriptor/se_t_tebd.py @@ -186,9 +186,10 @@ def _call_compressed( ) -> Any: """Compressed forward using tabulate_fusion_se_t_tebd custom op.""" # env_mat: nf x nloc x nnei x 4 + atype_ext_for_env = atype_ext.clamp_min(0) rr, _diff, sw = self.se_ttebd.env_mat.call( coord_ext, - atype_ext, + atype_ext_for_env, nlist, self.se_ttebd.mean[...], self.se_ttebd.stddev[...], diff --git a/source/tests/common/dpmodel/test_dpa1_call_graph_block.py b/source/tests/common/dpmodel/test_dpa1_call_graph_block.py index 11db9a7785..b60043f6fd 100644 --- a/source/tests/common/dpmodel/test_dpa1_call_graph_block.py +++ b/source/tests/common/dpmodel/test_dpa1_call_graph_block.py @@ -252,3 +252,29 @@ def test_strip_attn2_equals_dense(self, type_one_side) -> None: """attn_layer=2, smooth=False: bit-exact (avoids by-design smooth softmax divergence).""" dd = self._make(type_one_side, smooth=False, attn_layer=2) self._assert_parity(dd, compact=False) + + def test_virtual_center_uses_real_type_for_graph_statistics(self) -> None: + """Graph normalization clamps virtual centers before real-only stats.""" + dd = self._make(type_one_side=True, smooth=False, attn_layer=0) + dd.se_atten.mean[0, :, :] = 0.25 + dd.se_atten.mean[1, :, :] = -0.5 + coord = np.array([[[0.0, 0.0, 0.0], [1.0, 0.0, 0.0]]]) + nlist = np.array([[[1, -1], [0, -1]]], dtype=np.int64) + mapping = np.array([[0, 1]], dtype=np.int64) + graph = from_dense_quartet(coord, nlist, mapping, compact=False) + type_embedding = dd.type_embedding.call() + + actual, _ = dd.se_atten.call_graph( + graph, + np.array([-1, 1], dtype=np.int64), + type_embedding=type_embedding, + ) + expected, _ = dd.se_atten.call_graph( + graph, + np.array([0, 1], dtype=np.int64), + type_embedding=type_embedding, + ) + + self_value = actual[0] + assert np.isfinite(self_value).all() + np.testing.assert_allclose(self_value, expected[0]) diff --git a/source/tests/pt/model/test_virtual_type_embedding.py b/source/tests/pt/model/test_virtual_type_embedding.py index 70b6eee6b3..079467edad 100644 --- a/source/tests/pt/model/test_virtual_type_embedding.py +++ b/source/tests/pt/model/test_virtual_type_embedding.py @@ -57,6 +57,41 @@ def _assert_outputs_close(actual, expected) -> None: torch.testing.assert_close(actual_value, expected_value) +def _virtual_center_inputs(nnei: int): + device = env.DEVICE + coord = torch.tensor( + [[[0.0, 0.0, 0.0], [1.0, 0.0, 0.0], [0.0, 1.0, 0.0]]], + dtype=torch.float64, + device=device, + ) + nlist = torch.tensor([[[1, 2, -1, -1]]], dtype=torch.long, device=device)[ + :, :, :nnei + ] + virtual_atype = torch.tensor([[-1, 1, 0]], dtype=torch.long, device=device) + reference_atype = torch.tensor([[0, 1, 0]], dtype=torch.long, device=device) + type_embedding = torch.tensor( + [[1.0, 2.0], [3.0, 4.0], [0.0, 0.0]], + dtype=torch.float64, + device=device, + ) + virtual_embedding = type_embedding[ + torch.where( + virtual_atype >= 0, + virtual_atype, + torch.full_like(virtual_atype, 2), + ) + ] + return ( + coord, + nlist, + virtual_atype, + reference_atype, + type_embedding, + virtual_embedding, + type_embedding[reference_atype], + ) + + def test_se_atten_strip_virtual_neighbor_matches_padding() -> None: """Two-side DPA1 pair indices remap a virtual neighbor before folding.""" ( @@ -140,3 +175,92 @@ def test_se_t_tebd_strip_virtual_neighbor_matches_padding() -> None: ) _assert_outputs_close(actual, expected) + + +def test_se_atten_virtual_center_uses_real_type_statistics() -> None: + """DPA1 clamps a virtual center before indexing mean and stddev.""" + ( + coord, + nlist, + virtual_atype, + reference_atype, + type_embedding, + virtual_embedding, + reference_embedding, + ) = _virtual_center_inputs(4) + descriptor = DescrptBlockSeAtten( + rcut=4.0, + rcut_smth=0.5, + sel=[2, 2], + ntypes=2, + attn_layer=0, + axis_neuron=2, + neuron=[6, 12], + tebd_dim=2, + tebd_input_mode="strip", + type_one_side=True, + precision="float64", + seed=1, + ).to(env.DEVICE) + descriptor.mean[0, :, :] = 0.25 + descriptor.mean[1, :, :] = -0.5 + + actual = descriptor( + nlist, + coord, + virtual_atype, + virtual_embedding, + type_embedding=type_embedding, + ) + expected = descriptor( + nlist, + coord, + reference_atype, + reference_embedding, + type_embedding=type_embedding, + ) + + _assert_outputs_close(actual, expected) + + +def test_se_t_tebd_virtual_center_uses_real_type_statistics() -> None: + """SE_T clamps a virtual center before indexing mean and stddev.""" + ( + coord, + nlist, + virtual_atype, + reference_atype, + type_embedding, + virtual_embedding, + reference_embedding, + ) = _virtual_center_inputs(2) + descriptor = DescrptBlockSeTTebd( + rcut=4.0, + rcut_smth=0.5, + sel=2, + ntypes=2, + neuron=[4, 8], + tebd_dim=2, + tebd_input_mode="strip", + precision="float64", + seed=1, + ).to(env.DEVICE) + descriptor.mean[0, :, :] = 0.25 + descriptor.mean[1, :, :] = -0.5 + + actual = descriptor( + nlist, + coord, + virtual_atype, + virtual_embedding, + type_embedding=type_embedding, + ) + expected = descriptor( + nlist, + coord, + reference_atype, + reference_embedding, + type_embedding=type_embedding, + ) + + _assert_outputs_close(actual, expected) diff --git a/source/tests/pt_expt/descriptor/test_dpa1.py b/source/tests/pt_expt/descriptor/test_dpa1.py index a9ec8cde71..7b5070db3c 100644 --- a/source/tests/pt_expt/descriptor/test_dpa1.py +++ b/source/tests/pt_expt/descriptor/test_dpa1.py @@ -3,6 +3,9 @@ from types import ( SimpleNamespace, ) +from unittest.mock import ( + patch, +) import numpy as np import pytest @@ -14,6 +17,7 @@ from deepmd.dpmodel.descriptor.dpa1 import DescrptDPA1 as DPDescrptDPA1 from deepmd.pt_expt.descriptor.dpa1 import ( DescrptDPA1, + _env_mat, _strip_pair_index, ) from deepmd.pt_expt.utils import ( @@ -99,6 +103,48 @@ def test_strip_pair_index_remaps_virtual_types() -> None: ) +def test_fused_env_prologue_clamps_virtual_center_statistics() -> None: + """Compressed/Triton routes sanitize centers before real-only stats.""" + device = env.DEVICE + descriptor = DescrptDPA1( + rcut=4.0, + rcut_smth=0.5, + sel=[2, 2], + ntypes=2, + attn_layer=0, + axis_neuron=2, + neuron=[6, 12], + tebd_dim=2, + tebd_input_mode="strip", + type_one_side=True, + seed=GLOBAL_SEED, + ).to(device) + descriptor.se_atten.mean[0, :, :] = 0.25 + descriptor.se_atten.mean[1, :, :] = -0.5 + coord = torch.tensor( + [[[0.0, 0.0, 0.0], [1.0, 0.0, 0.0], [0.0, 1.0, 0.0]]], + dtype=torch.float64, + device=device, + ) + nlist = torch.tensor([[[1, 2, -1, -1]]], dtype=torch.long, device=device) + + with patch("deepmd.pt_expt.descriptor.dpa1.triton_infer_level", return_value=0): + actual = _env_mat( + descriptor, + coord, + torch.tensor([[-1, 1, 0]], dtype=torch.long, device=device), + nlist, + )[5] + expected = _env_mat( + descriptor, + coord, + torch.tensor([[0, 1, 0]], dtype=torch.long, device=device), + nlist, + )[5] + + torch.testing.assert_close(actual, expected) + + class TestDescrptDPA1(TestCaseSingleFrameWithNlist): def setup_method(self) -> None: TestCaseSingleFrameWithNlist.setUp(self) From 609d77f2d9a532d11dc2db88fa34df6d0ef0e31b Mon Sep 17 00:00:00 2001 From: njzjz-bot Date: Sun, 2 Aug 2026 00:53:08 +0800 Subject: [PATCH 5/7] test(pt_expt): cover fused virtual center statistics Coding-Agent: Codex\nCodex-Version: codex-cli 0.144.6\nModel: gpt-5.6-sol\nReasoning-Effort: xhigh --- source/tests/pt_expt/descriptor/test_dpa1.py | 42 ++++++++++++++++---- 1 file changed, 35 insertions(+), 7 deletions(-) diff --git a/source/tests/pt_expt/descriptor/test_dpa1.py b/source/tests/pt_expt/descriptor/test_dpa1.py index 7b5070db3c..47c911a834 100644 --- a/source/tests/pt_expt/descriptor/test_dpa1.py +++ b/source/tests/pt_expt/descriptor/test_dpa1.py @@ -127,21 +127,49 @@ def test_fused_env_prologue_clamps_virtual_center_statistics() -> None: device=device, ) nlist = torch.tensor([[[1, 2, -1, -1]]], dtype=torch.long, device=device) + virtual_atype = torch.tensor([[-1, 1, 0]], dtype=torch.long, device=device) + real_atype = torch.tensor([[0, 1, 0]], dtype=torch.long, device=device) with patch("deepmd.pt_expt.descriptor.dpa1.triton_infer_level", return_value=0): - actual = _env_mat( - descriptor, - coord, - torch.tensor([[-1, 1, 0]], dtype=torch.long, device=device), - nlist, - )[5] expected = _env_mat( descriptor, coord, - torch.tensor([[0, 1, 0]], dtype=torch.long, device=device), + real_atype, nlist, )[5] + captured_center_types = [] + + def fused_env_mat_stub( + coord_ext, + nlist_arg, + center_types, + mean, + stddev, + rcut, + rcut_smth, + **kwargs, + ): + """Mirror the dense result while exposing the fused center-type input.""" + captured_center_types.append(center_types.clone()) + return descriptor.se_atten.env_mat.call( + coord_ext, real_atype, nlist_arg, mean, stddev + ) + + with ( + patch("deepmd.pt_expt.descriptor.dpa1.triton_infer_level", return_value=1), + patch( + "deepmd.pt_expt.descriptor.dpa1._env_mat_triton", + side_effect=fused_env_mat_stub, + ), + ): + actual = _env_mat(descriptor, coord, virtual_atype, nlist)[5] + + assert len(captured_center_types) == 1 + torch.testing.assert_close( + captured_center_types[0], + torch.zeros((1, 1), dtype=torch.long, device=device), + ) torch.testing.assert_close(actual, expected) From 6d57520b77454a284e4e7f5d22b18e0ff8ab5183 Mon Sep 17 00:00:00 2001 From: njzjz-bot Date: Sun, 2 Aug 2026 22:04:29 +0800 Subject: [PATCH 6/7] test(dpmodel): make DPA3 virtual-type regression fail pre-fix A virtual -1 sentinel wraps to the final padding row under NumPy and array_api_strict, so the old DPA3 test passed even without the remap. Use -2 as the negative sentinel so a raw take lands on a real row and only the explicit remap-to-padding satisfies the test. Coding-Agent: opencode opencode-Version: 1.18.9 Model: ustc/deepseek-v4-flash Reasoning-Effort: max --- .../dpmodel/test_type_embedding_virtual.py | 73 ++++++++++++------- 1 file changed, 47 insertions(+), 26 deletions(-) diff --git a/source/tests/common/dpmodel/test_type_embedding_virtual.py b/source/tests/common/dpmodel/test_type_embedding_virtual.py index b809830fec..6e2d39079d 100644 --- a/source/tests/common/dpmodel/test_type_embedding_virtual.py +++ b/source/tests/common/dpmodel/test_type_embedding_virtual.py @@ -4,6 +4,9 @@ import array_api_strict import numpy as np import pytest +from typing import ( + Any, +) from deepmd.dpmodel.common import ( to_numpy_array, @@ -215,25 +218,34 @@ def test_dpa2_virtual_neighbor_matches_explicit_padding() -> None: def test_dpa3_virtual_type_uses_padding_for_both_mapping_paths( use_loc_mapping: bool, ) -> None: - """DPA3 remaps before either local-only or full extended gathering.""" - descriptor = DescrptDPA3( - ntypes=2, - repflow=RepFlowArgs( - n_dim=8, - e_dim=4, - a_dim=4, - nlayers=1, - e_rcut=4.0, - e_rcut_smth=0.5, - e_sel=4, - a_rcut=4.0, - a_rcut_smth=0.5, - a_sel=3, - axis_neuron=2, - update_angle=False, - ), - use_loc_mapping=use_loc_mapping, - seed=17, + """DPA3 remaps before either local-only or full extended gathering. + + A virtual ``-1`` is deliberately avoided for the sentinel: under both + NumPy and array_api_strict, a raw ``take`` wraps ``-1`` to the final + (padding) row, which would coincide with a correct remap and let a + regression slip through. ``-2`` wraps to a real row instead, so the test + only passes when the negative type is explicitly remapped to padding. + """ + descriptor = convert_array_api_strict_value( + DescrptDPA3( + ntypes=2, + repflow=RepFlowArgs( + n_dim=8, + e_dim=4, + a_dim=4, + nlayers=1, + e_rcut=4.0, + e_rcut_smth=0.5, + e_sel=4, + a_rcut=4.0, + a_rcut_smth=0.5, + a_sel=3, + axis_neuron=2, + update_angle=False, + ), + use_loc_mapping=use_loc_mapping, + seed=17, + ) ) table = np.vstack( ( @@ -242,12 +254,13 @@ def test_dpa3_virtual_type_uses_padding_for_both_mapping_paths( np.zeros(descriptor.tebd_dim), ) ) - descriptor.type_embedding.call = lambda: table + strict_table = array_api_strict.asarray(table) + descriptor.type_embedding.call = lambda: strict_table class CaptureRepflows: """Capture DPA3's initial node embeddings without later env lookups.""" - node_ebd_ext: np.ndarray + node_ebd_ext: Any def __call__( self, @@ -270,12 +283,20 @@ def __call__( capture = CaptureRepflows() descriptor.repflows = capture - coord = np.array([[[0.0, 0.0, 0.0], [1.0, 0.0, 0.0], [0.0, 1.0, 0.0]]]) - nlist = np.array([[[1, 2, -1, -1], [0, 2, -1, -1]]], dtype=np.int64) - mapping = np.array([[0, 1, 0]], dtype=np.int64) - atype = np.array([[0, -1, 1] if use_loc_mapping else [0, 1, -1]], dtype=np.int64) + coord = array_api_strict.asarray( + np.array([[[0.0, 0.0, 0.0], [1.0, 0.0, 0.0], [0.0, 1.0, 0.0]]]) + ) + nlist = array_api_strict.asarray( + np.array([[[1, 2, -1, -1], [0, 2, -1, -1]]], dtype=np.int64) + ) + mapping = array_api_strict.asarray(np.array([[0, 1, 0]], dtype=np.int64)) + atype = array_api_strict.asarray( + np.array([[0, -2, 1] if use_loc_mapping else [0, 1, -2]], dtype=np.int64) + ) descriptor(coord, atype, nlist, mapping) expected_rows = [0, 2] if use_loc_mapping else [0, 1, 2] - np.testing.assert_array_equal(capture.node_ebd_ext, table[expected_rows][None, ...]) + np.testing.assert_array_equal( + to_numpy_array(capture.node_ebd_ext), table[expected_rows][None, ...] + ) From 078ea56a8103fbc63b401f1168ffec943fb7419a Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sun, 2 Aug 2026 14:05:25 +0000 Subject: [PATCH 7/7] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- source/tests/common/dpmodel/test_type_embedding_virtual.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/source/tests/common/dpmodel/test_type_embedding_virtual.py b/source/tests/common/dpmodel/test_type_embedding_virtual.py index 6e2d39079d..829eff0881 100644 --- a/source/tests/common/dpmodel/test_type_embedding_virtual.py +++ b/source/tests/common/dpmodel/test_type_embedding_virtual.py @@ -1,13 +1,14 @@ # SPDX-License-Identifier: LGPL-3.0-or-later """Regression tests for virtual types in dpmodel embedding gathers.""" -import array_api_strict -import numpy as np -import pytest from typing import ( Any, ) +import array_api_strict +import numpy as np +import pytest + from deepmd.dpmodel.common import ( to_numpy_array, )