deepmd.dpmodel.descriptor.unimol_nn.encoder

Uni-Mol v1 transformer encoder, ported to the array-API dpmodel layer.

The maths follows Uni-Mol (deepmodeling/Uni-Mol) at commit 90f52c4 and the Uni-Core modules it builds on (dptech-corp/Uni-Core) at commit ace6fae, both MIT licensed:

Copyright (c) DP Technology This source code is licensed under the MIT license found in the LICENSE file in the root directory of that source tree.

Parts of Uni-Core derive in turn from fairseq (Copyright (c) Facebook, Inc. and its affiliates, MIT licensed).

Ported files and the classes taken from them:

The port keeps upstream’s exact order of operations, so that a checkpoint converted from Uni-Mol reproduces upstream outputs to floating-point rounding. Two upstream quirks are reproduced deliberately and are marked in the code: the truncated pi in the Gaussian basis, and the fp32 cast inside the norm regularisers.

Attributes

UNIMOL_PI

Classes

NonLinearHead	Two-layer head, linear1 -> activation -> linear2.
GaussianLayer	Gaussian radial basis with a per-edge-type affine on the distance.
SelfMultiheadAttention	Self-attention with a fused QKV projection and an additive pair bias.
TransformerEncoderLayer	Pre-layer-norm transformer block that also returns attention logits.
TransformerEncoderWithPair	Stack of pre-LN blocks that carries a pair representation.

Functions

softmax(x[, axis])	Numerically stable softmax over axis.
dropout(x, p, training)	Apply dropout, but only while training.
norm_loss(x[, eps, tolerance])	Hinge on the deviation of the row norm from sqrt(dim).
masked_mean(mask, value[, axis, eps])	Mean of value over mask, then mean over what is left.

Module Contents

deepmd.dpmodel.descriptor.unimol_nn.encoder.UNIMOL_PI = 3.14159

deepmd.dpmodel.descriptor.unimol_nn.encoder.softmax(x, axis: int = -1)

Numerically stable softmax over axis.

Rows that are entirely -inf would produce NaN, as they do upstream. Uni-Mol never builds such a row: the BOS column is never masked.

deepmd.dpmodel.descriptor.unimol_nn.encoder.dropout(x, p: float, training: bool)

Apply dropout, but only while training.

deepmd has no dropout anywhere else, and the array API has no random numbers, so this dispatches to torch when a training step actually needs it. Inference, which is what the array-API backends are for, is the identity. Training on a non-torch backend is refused rather than silently dropping the regularisation, which would be a quiet parity bug.

deepmd.dpmodel.descriptor.unimol_nn.encoder.norm_loss(x, eps: float = 1e-10, tolerance: float = 1.0)

Hinge on the deviation of the row norm from sqrt(dim).

Mirrors norm_loss in unimol/models/transformer_encoder_with_pair.py:101. Upstream evaluates this in fp32 because it pretrains a pure-fp16 model; the cast is reproduced so that the value matches upstream exactly.

deepmd.dpmodel.descriptor.unimol_nn.encoder.masked_mean(mask, value, axis=-1, eps: float = 1e-10)

Mean of value over mask, then mean over what is left.

Mirrors masked_mean in transformer_encoder_with_pair.py:108. The eps in the denominator is what makes an all-padding row return 0 rather than NaN.

class deepmd.dpmodel.descriptor.unimol_nn.encoder.NonLinearHead(input_dim: int, out_dim: int, activation_function: str = 'gelu_erf', hidden: int | None = None, precision: str = 'float64', seed: int | list[int] | None = None)

Bases: deepmd.dpmodel.common.NativeOP

Two-layer head, linear1 -> activation -> linear2.

Mirrors unimol/models/unimol.py:347.

input_dim

out_dim

activation_function = 'gelu_erf'

hidden

precision = 'float64'

linear1

linear2

call(x)

Forward pass in NumPy implementation.

serialize() → dict

classmethod deserialize(data: dict) → NonLinearHead

class deepmd.dpmodel.descriptor.unimol_nn.encoder.GaussianLayer(k: int = 128, edge_types: int = 1024, single_precision_basis: bool = True, precision: str = 'float64', seed: int | list[int] | None = None)

Bases: deepmd.dpmodel.common.NativeOP

Gaussian radial basis with a per-edge-type affine on the distance.

Mirrors unimol/models/unimol.py:400. means and stds are shared by all edge types; mul and bias hold one scalar per ordered element pair.

k = 128

edge_types = 1024

single_precision_basis = True

precision = 'float64'

means

stds

mul

bias

call(dist, edge_type)

Expand dist (nf x nt x nt) into k Gaussians per atom pair.

serialize() → dict

Serialize the basis.

classmethod deserialize(data: dict) → GaussianLayer

Deserialize the basis.

class deepmd.dpmodel.descriptor.unimol_nn.encoder.SelfMultiheadAttention(embed_dim: int, num_heads: int, dropout: float = 0.1, bias: bool = True, scaling_factor: float = 1.0, precision: str = 'float64', seed: int | list[int] | None = None)

Bases: deepmd.dpmodel.common.NativeOP

Self-attention with a fused QKV projection and an additive pair bias.

Mirrors unicore/modules/multihead_attention.py:12. Only the return_attn path of upstream is kept, because Uni-Mol always takes it: the pre-softmax logits are the pair representation that the next layer biases with.

embed_dim

num_heads

dropout = 0.1

head_dim

scaling

precision = 'float64'

in_proj

out_proj

call(query, attn_bias, training: bool = False)

Return the attended values and the biased pre-softmax logits.

Parameters:

query

nf x nt x embed_dim.

attn_bias

(nf * num_heads) x nt x nt, already carrying -inf on padded keys.

Returns:

tuple

(output, attn_weights), shaped nf x nt x embed_dim and (nf * num_heads) x nt x nt.

serialize() → dict

classmethod deserialize(data: dict) → SelfMultiheadAttention

class deepmd.dpmodel.descriptor.unimol_nn.encoder.TransformerEncoderLayer(embed_dim: int = 768, ffn_embed_dim: int = 3072, attention_heads: int = 8, dropout: float = 0.1, attention_dropout: float = 0.1, activation_dropout: float = 0.0, activation_function: str = 'gelu_erf', precision: str = 'float64', seed: int | list[int] | None = None)

Bases: deepmd.dpmodel.common.NativeOP

Pre-layer-norm transformer block that also returns attention logits.

Mirrors unicore/modules/transformer_encoder_layer.py:15. Uni-Mol never sets post_ln, so only the pre-LN order is implemented.

embed_dim = 768

ffn_embed_dim = 3072

attention_heads = 8

dropout = 0.1

attention_dropout = 0.1

activation_dropout = 0.0

activation_function = 'gelu_erf'

precision = 'float64'

self_attn

self_attn_layer_norm

fc1

fc2

final_layer_norm

call(x, attn_bias, training: bool = False)

Forward pass in NumPy implementation.

serialize() → dict

classmethod deserialize(data: dict) → TransformerEncoderLayer

class deepmd.dpmodel.descriptor.unimol_nn.encoder.TransformerEncoderWithPair(encoder_layers: int = 6, embed_dim: int = 768, ffn_embed_dim: int = 3072, attention_heads: int = 8, emb_dropout: float = 0.1, dropout: float = 0.1, attention_dropout: float = 0.1, activation_dropout: float = 0.0, max_seq_len: int = 256, activation_function: str = 'gelu_erf', no_final_head_layer_norm: bool = False, precision: str = 'float64', seed: int | list[int] | None = None)

Bases: deepmd.dpmodel.common.NativeOP

Stack of pre-LN blocks that carries a pair representation.

Mirrors unimol/models/transformer_encoder_with_pair.py:14. Each block’s pre-softmax logits become the next block’s bias, so the pair representation is the running sum of per-layer logits. Returns the node representation, the final pair representation, the pair delta, and the two norm regularisers.

encoder_layers = 6

embed_dim = 768

ffn_embed_dim = 3072

attention_heads = 8

emb_dropout = 0.1

dropout = 0.1

attention_dropout = 0.1

activation_dropout = 0.0

max_seq_len = 256

activation_function = 'gelu_erf'

no_final_head_layer_norm = False

precision = 'float64'

emb_layer_norm

final_layer_norm

final_head_layer_norm = None

layers

call(emb, attn_mask, padding_mask, training: bool = False)

Run the stack.

Parameters:

emb

nf x nt x embed_dim token embedding.

attn_mask

(nf * heads) x nt x nt bias from the Gaussian basis.

padding_mask

nf x nt, 1 on padded positions.

Returns:

tuple

(x, pair_rep, delta_pair_rep, x_norm, delta_pair_rep_norm).

serialize() → dict

classmethod deserialize(data: dict) → TransformerEncoderWithPair