This issue comes from a Codex global scan of deepmodeling/deepmd-kit at commit 73de44b1f94471b2e3bdb6b11f57b34d7bc791bb.
Problem
Several dpmodel descriptors allocate padding type-embedding rows but still gather with raw atom types. Negative placeholder types can therefore either crash on backends that reject negative gather indices, or wrap to the last row on NumPy-like backends.
Examples:
|
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), |
|
(nf, nall, self.tebd_dim), |
|
) |
|
xp = array_api_compat.array_namespace(graph.edge_vec) |
|
dev = array_api_compat.device(graph.edge_vec) |
|
grrg, rot_mat = self.se_atten.call_graph( |
|
graph, atype, type_embedding=type_embedding |
|
) |
|
# FLAT node axis (N, ...): no (nf, nloc) reshape -- ragged-native, spec. |
|
if self.concat_output_tebd: |
|
tebd = xp.asarray(type_embedding, device=dev) |
|
atype_local = xp.asarray(atype, device=dev) |
|
atype_embd = xp.take(tebd, atype_local, axis=0) # (N, tebd_dim) |
|
grrg = xp.concat([grrg, atype_embd], axis=-1) |
|
type_embedding = self.type_embedding.call() |
|
# repinit |
|
g1_ext = xp.reshape( |
|
xp.take(type_embedding, xp.reshape(atype_ext, (-1,)), axis=0), |
|
(nframes, nall, self.tebd_dim), |
|
) |
|
g1_inp = xp_take_first_n(g1_ext, 1, nloc) |
|
type_embedding = self.type_embedding.call() |
|
if self.use_loc_mapping: |
|
node_ebd_ext = xp.reshape( |
|
xp.take( |
|
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), |
|
(nframes, nall, self.tebd_dim), |
|
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), |
|
(nf, nall, self.tebd_dim), |
|
) |
|
# nfnl x tebd_dim |
|
atype_embd = atype_embd_ext[:, :nloc, :] |
|
# === Step 2. Type embedding (l=0) === |
|
# Use ``xp_take_first_n`` (torch.index_select) rather than a plain |
|
# ``[:, :nloc]`` slice: the slice makes torch.export emit a spurious |
|
# ``Ne(nall, nloc)`` contiguity guard that breaks the ``nall == nloc`` |
|
# (NoPBC, no ghost atoms) case in the compiled .pt2 artifact. |
|
atype_loc = xp_take_first_n(atype_ext, 1, nloc) |
|
type_ebed = xp.reshape( |
|
self.type_embedding(atype_loc), (n_nodes, self.channels) |
|
) # (N, C) |
SeZMTypeEmbedding can create an explicit padding row:
|
# from pt's torch generator; weight values are not bit-compatible. |
|
init_std = 1.0 / math.sqrt(float(self.ntypes + self.embed_dim)) |
|
rng = np.random.default_rng(child_seed(seed, 0)) |
|
table = rng.normal(scale=init_std, size=(self.ntypes, self.embed_dim)) |
|
if self.padding: |
|
table = np.concatenate( |
|
[table, np.zeros((1, self.embed_dim), dtype=table.dtype)], axis=0 |
|
) |
|
self.adam_type_embedding = table.astype(prec) |
But the gather contract says negative type ids are invalid and the implementation passes raw indices to xp.take:
|
def call(self, atype: Any) -> Any: |
|
""" |
|
Gather type embeddings. |
|
|
|
Parameters |
|
---------- |
|
atype |
|
Atom types with shape (...,). Valid type range is [0, ntypes-1] |
|
(plus the padding row index ``ntypes`` when ``padding=True``). |
|
Negative type ids are invalid input and are NOT validated here |
|
(caller contract). |
|
|
|
Returns |
|
------- |
|
Array |
|
Type embeddings with shape (..., embed_dim). |
|
""" |
|
xp = array_api_compat.array_namespace(atype) |
|
weight = xp_asarray_nodetach( |
|
xp, self.adam_type_embedding[...], device=array_api_compat.device(atype) |
|
) |
|
# pt embedding.py:143 torch.embedding -> flat int64 take + reshape. |
|
index = xp.astype(xp.reshape(atype, (-1,)), xp.int64) |
|
out = xp.take(weight, index, axis=0) |
|
return xp.reshape(out, (*atype.shape, self.embed_dim)) |
The environment-seed embedding path also reuses raw source and destination atom types:
|
src_index = xp.astype(xp.reshape(src, (n_edge,)), xp.int64) |
|
dst_index = xp.astype(xp.reshape(dst, (n_edge,)), xp.int64) |
|
atype_src = xp.take(atype_flat, src_index, axis=0) # (E,) |
|
atype_dst = xp.take(atype_flat, dst_index, axis=0) # (E,) |
|
type_src = self.env_type_embed(atype_src) # (E, type_dim) |
|
type_dst = self.env_type_embed(atype_dst) # (E, type_dim) |
This is distinct from #5628, which covers env-mat normalization for virtual center atoms. This issue is about descriptor type-embedding gathers using raw negative atom-type ids.
Impact
Mixed-size batches or graph paths that include atype == -1 placeholders can fail or produce backend-dependent descriptor embeddings instead of using the intended padding row.
Suggested fix
Before every type-embedding gather, remap negative atom types to the explicit padding row index, usually ntypes, when a padding row exists. Otherwise validate and reject negative atom types before gathering. Add backend parity tests with atype_ext containing -1.
This issue comes from a Codex global scan of
deepmodeling/deepmd-kitat commit73de44b1f94471b2e3bdb6b11f57b34d7bc791bb.Problem
Several dpmodel descriptors allocate padding type-embedding rows but still gather with raw atom types. Negative placeholder types can therefore either crash on backends that reject negative gather indices, or wrap to the last row on NumPy-like backends.
Examples:
deepmd-kit/deepmd/dpmodel/descriptor/dpa1.py
Lines 702 to 707 in 73de44b
deepmd-kit/deepmd/dpmodel/descriptor/dpa1.py
Lines 760 to 770 in 73de44b
deepmd-kit/deepmd/dpmodel/descriptor/dpa2.py
Lines 893 to 899 in 73de44b
deepmd-kit/deepmd/dpmodel/descriptor/dpa3.py
Lines 707 to 720 in 73de44b
deepmd-kit/deepmd/dpmodel/descriptor/se_t_tebd.py
Lines 398 to 405 in 73de44b
deepmd-kit/deepmd/dpmodel/descriptor/dpa4.py
Lines 838 to 846 in 73de44b
SeZMTypeEmbeddingcan create an explicit padding row:deepmd-kit/deepmd/dpmodel/descriptor/dpa4_nn/embedding.py
Lines 135 to 143 in 73de44b
But the gather contract says negative type ids are invalid and the implementation passes raw indices to
xp.take:deepmd-kit/deepmd/dpmodel/descriptor/dpa4_nn/embedding.py
Lines 145 to 169 in 73de44b
The environment-seed embedding path also reuses raw source and destination atom types:
deepmd-kit/deepmd/dpmodel/descriptor/dpa4_nn/embedding.py
Lines 617 to 622 in 73de44b
This is distinct from #5628, which covers env-mat normalization for virtual center atoms. This issue is about descriptor type-embedding gathers using raw negative atom-type ids.
Impact
Mixed-size batches or graph paths that include
atype == -1placeholders can fail or produce backend-dependent descriptor embeddings instead of using the intended padding row.Suggested fix
Before every type-embedding gather, remap negative atom types to the explicit padding row index, usually
ntypes, when a padding row exists. Otherwise validate and reject negative atom types before gathering. Add backend parity tests withatype_extcontaining-1.