fix(dpmodel): remap virtual type embeddings (deepmodeling#5856)

Closes deepmodeling#5665.

## Summary

- centralize remapping of negative virtual atom types to the reserved
final row of padded embedding tables
- apply the same pair-index convention across dpmodel, pt_expt, legacy
PyTorch, and Paddle DPA1/DPA2/SE_T implementations
- keep real-type-only exclusion and normalization tables on a separate
clamp-to-real convention
- use the shared remapping helper in SeZMTypeEmbedding
- cover DPA1 virtual neighbors inside the cutoff, DPA2, both DPA3
local-mapping branches, pt_expt pair arithmetic, and legacy PyTorch
descriptors

## Sentinel convention

Only tables explicitly constructed with a final padding row may use
ntypes as the remapped virtual type. Real-type-only tables such as davg,
dstd, exclusion inputs, and spin masks must instead receive masked or
clamped real-type indices. TypeEmbedNet reconstructs a literal zero
padding row; SeZMTypeEmbedding reserves and initializes its stored final
row to zero.

## Validation

- focused virtual-type and backend regressions: 15 passed
- DPA1 graph parity plus the existing DPA2 descriptor tests: 18 passed
- pt_expt DPA1 export tests: 8 passed
- changed-file ruff check: passed
- ruff format .: passed
- git diff --check: passed
- Paddle-specific code was statically checked but could not be executed
locally because Paddle is not installed
- full ruff check . still reports five unrelated pre-existing findings
in deepmd/jax/jax_md/__init__.py and deepmd/tf/entrypoints/__init__.py

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

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

- **New Features**
- Added consistent padding-aware type-embedding handling across dense,
graph, and accelerated descriptor calculations.
- Virtual or negative atom types now map safely to the designated
padding embedding.

- **Bug Fixes**
- Prevented incorrect embedding lookups and invalid type indexing across
supported execution paths.
- Ensured virtual types produce results equivalent to explicit padding
types.

- **Tests**
- Added regression coverage across NumPy, PyTorch, descriptor modes, and
embedding pathways.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Co-authored-by: njzjz-bot <njzjz.bot@gmail.com>
Co-authored-by: njzjz-bot <njzjz-bot@users.noreply.github.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: njzjz-bot

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,
@@ -904,7 +906,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
@@ -1006,7 +1008,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)
@@ -1839,6 +1841,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:
@@ -1853,9 +1856,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,),
                 )
@@ -2016,22 +2021,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,
@@ -2061,9 +2071,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)
@@ -2147,6 +2157,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:

deepmd/dpmodel/descriptor/dpa2.py

@@ -39,6 +39,7 @@
 )
 from deepmd.dpmodel.utils.type_embed import (
     TypeEmbedNet,
+    take_type_embedding,
 )
 from deepmd.dpmodel.utils.update_sel import (
     UpdateSel,
@@ -1349,7 +1350,7 @@ def _call_dense(
         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)

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,
@@ -740,16 +741,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),
             )
 

deepmd/dpmodel/descriptor/dpa4_nn/embedding.py

@@ -42,6 +42,9 @@
 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,
 )
@@ -146,6 +149,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 = remap_atype_to_padding(index, self.ntypes + 1)
         out = xp.take(weight, index, axis=0)
         return xp.reshape(out, (*atype.shape, self.embed_dim))
 

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
@@ -925,6 +927,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))

deepmd/dpmodel/utils/type_embed.py

@@ -31,6 +31,75 @@ 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 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,
+        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.
+
+    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
+    # 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.
 

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,
@@ -582,6 +583,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 +599,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

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,
@@ -929,6 +930,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])

deepmd/pt/model/descriptor/se_atten.py

@@ -769,12 +769,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,
@@ -897,6 +898,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:
@@ -906,8 +915,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: