deepmd.pt_expt.descriptor.dpa1

Attributes

_DESTINATION_CSR_FIELDS
_DUAL_CSR_FIELDS

Classes

DescrptDPA1

Attention-based descriptor which is proposed in the pretrainable DPA-1[1] model.

Functions

_has_graph_fields(→ bool)
_env_mat(→ tuple)	Environment-matrix prologue shared by every fused path (strip / concat).
_strip_pair_index(→ torch.Tensor)	Strip-mode per-edge row index into the type-pair embedding table.
_concat_embedding_input(→ torch.Tensor)	Concat-mode embedding input: radial channel concatenated with the
_grrg_from_moment(→ Any)	Strip-mode epilogue: symmetry-invariant contraction of the moment.
_type_pair_table(→ torch.Tensor)	Type-pair embedding table (P, ng) gathered per edge by tebd_idx.
_without_magnetic_force(→ tuple[torch.Tensor, ...])	Append the empty magnetic force of a spin-free fused composition.

Module Contents

deepmd.pt_expt.descriptor.dpa1._has_graph_fields(graph: Any, fields: tuple[str, ...]) → bool

deepmd.pt_expt.descriptor.dpa1._DESTINATION_CSR_FIELDS = ('destination_order', 'destination_row_ptr')

deepmd.pt_expt.descriptor.dpa1._DUAL_CSR_FIELDS = ('destination_order', 'destination_row_ptr', 'source_order', 'source_row_ptr')

deepmd.pt_expt.descriptor.dpa1._env_mat(desc: Any, coord_ext: torch.Tensor, atype_ext: torch.Tensor, nlist: torch.Tensor) → tuple

Environment-matrix prologue shared by every fused path (strip / concat).

Returns (nf, nloc, nnei, ng, nfnl, rr, moment_basis, ss, sw, nlist_masked, type_embedding). rr is the four-component environment matrix, while moment_basis has four or nine components according to lmax. Excluded and padding edges are zeroed in both tensors.

deepmd.pt_expt.descriptor.dpa1._strip_pair_index(desc: Any, atype_ext: torch.Tensor, nlist_masked: torch.Tensor, type_embedding: torch.Tensor, nf: int, nloc: int, nnei: int) → torch.Tensor

Strip-mode per-edge row index into the type-pair embedding table.

One-side uses the neighbor type; two-side folds the (center, neighbor) pair into a flat index, mirroring the dense reference. Shape (nfnl*nnei,).

deepmd.pt_expt.descriptor.dpa1._concat_embedding_input(desc: Any, ss: torch.Tensor, atype_ext: torch.Tensor, nlist_masked: torch.Tensor, type_embedding: torch.Tensor, nf: int, nloc: int, nnei: int) → torch.Tensor

Concat-mode embedding input: radial channel concatenated with the neighbor (and, two-side, center) type embeddings.

Returns (nfnl, nnei, 1 + k * tebd_dim) with k in {1, 2} (one-side / two-side), matching the dense concat reference exactly.

deepmd.pt_expt.descriptor.dpa1._grrg_from_moment(desc: Any, xyz_scatter: torch.Tensor, type_embedding: torch.Tensor, coord_ext: torch.Tensor, atype_ext: torch.Tensor, nf: int, nloc: int, nnei: int, ng: int, sw: torch.Tensor) → Any

Strip-mode epilogue: symmetry-invariant contraction of the moment.

Consumes the unnormalized moment, applies the 1 / nnei normalization, forms the G^T G descriptor and the rotation matrix, and appends the center type embedding when concat_output_tebd.

deepmd.pt_expt.descriptor.dpa1._type_pair_table(desc: Any, type_embedding: torch.Tensor) → torch.Tensor

Type-pair embedding table (P, ng) gathered per edge by tebd_idx.

One-side keeps the neighbor-type rows; two-side forms every (neighbor, center) pair, mirroring the dense reference.

deepmd.pt_expt.descriptor.dpa1._without_magnetic_force(output: tuple[torch.Tensor, ...]) → tuple[torch.Tensor, ...]

Append the empty magnetic force of a spin-free fused composition.

Every implementation of fused_energy_force_graph returns the same six outputs, so the caller reads the magnetic force by position rather than by probing the arity of the result.

Parameters:

output

The five spin-free outputs: energy, atom energy, force, virial and atom virial.

Returns:

tuple of torch.Tensor

The same outputs followed by an empty magnetic force.

class deepmd.pt_expt.descriptor.dpa1.DescrptDPA1(*args: Any, **kwargs: Any)

Bases: deepmd.dpmodel.descriptor.dpa1.DescrptDPA1

Attention-based descriptor which is proposed in the pretrainable DPA-1[1] model.

This descriptor, \(\mathcal{D}^i \in \mathbb{R}^{M \times M_{<}}\), is given by

\[\mathcal{D}^i = \frac{1}{N_c^2}(\hat{\mathcal{G}}^i)^T \mathcal{R}^i (\mathcal{R}^i)^T \hat{\mathcal{G}}^i_<,\]

where \(\hat{\mathcal{G}}^i\) represents the embedding matrix:math:mathcal{G}^i after additional self-attention mechanism and \(\mathcal{R}^i\) is defined by the full case in the se_e2_a descriptor. Note that we obtain \(\mathcal{G}^i\) using the type embedding method by default in this descriptor.

To perform the self-attention mechanism, the queries \(\mathcal{Q}^{i,l} \in \mathbb{R}^{N_c\times d_k}\), keys \(\mathcal{K}^{i,l} \in \mathbb{R}^{N_c\times d_k}\), and values \(\mathcal{V}^{i,l} \in \mathbb{R}^{N_c\times d_v}\) are first obtained:

\[\left(\mathcal{Q}^{i,l}\right)_{j}=Q_{l}\left(\left(\mathcal{G}^{i,l-1}\right)_{j}\right),\]

\[\left(\mathcal{K}^{i,l}\right)_{j}=K_{l}\left(\left(\mathcal{G}^{i,l-1}\right)_{j}\right),\]

\[\left(\mathcal{V}^{i,l}\right)_{j}=V_{l}\left(\left(\mathcal{G}^{i,l-1}\right)_{j}\right),\]

where \(Q_{l}\), \(K_{l}\), \(V_{l}\) represent three trainable linear transformations that output the queries and keys of dimension \(d_k\) and values of dimension \(d_v\), and \(l\) is the index of the attention layer. The input embedding matrix to the attention layers, denoted by \(\mathcal{G}^{i,0}\), is chosen as the two-body embedding matrix.

Then the scaled dot-product attention method is adopted:

\[A(\mathcal{Q}^{i,l}, \mathcal{K}^{i,l}, \mathcal{V}^{i,l}, \mathcal{R}^{i,l})=\varphi\left(\mathcal{Q}^{i,l}, \mathcal{K}^{i,l},\mathcal{R}^{i,l}\right)\mathcal{V}^{i,l},\]

where \(\varphi\left(\mathcal{Q}^{i,l}, \mathcal{K}^{i,l},\mathcal{R}^{i,l}\right) \in \mathbb{R}^{N_c\times N_c}\) is attention weights. In the original attention method, one typically has \(\varphi\left(\mathcal{Q}^{i,l}, \mathcal{K}^{i,l}\right)=\mathrm{softmax}\left(\frac{\mathcal{Q}^{i,l} (\mathcal{K}^{i,l})^{T}}{\sqrt{d_{k}}}\right)\), with \(\sqrt{d_{k}}\) being the normalization temperature. This is slightly modified to incorporate the angular information:

\[\varphi\left(\mathcal{Q}^{i,l}, \mathcal{K}^{i,l},\mathcal{R}^{i,l}\right) = \mathrm{softmax}\left(\frac{\mathcal{Q}^{i,l} (\mathcal{K}^{i,l})^{T}}{\sqrt{d_{k}}}\right) \odot \hat{\mathcal{R}}^{i}(\hat{\mathcal{R}}^{i})^{T},\]

where \(\hat{\mathcal{R}}^{i} \in \mathbb{R}^{N_c\times 3}\) denotes normalized relative coordinates,

\(\hat{\mathcal{R}}^{i}_{j} = \frac{\boldsymbol{r}_{ij}}{\lVert \boldsymbol{r}_{ij} \lVert}\) and \(\odot\) means element-wise multiplication.

Then layer normalization is added in a residual way to finally obtain the self-attention local embedding matrix

\(\hat{\mathcal{G}}^{i} = \mathcal{G}^{i,L_a}\) after \(L_a\) attention layers:[^1]

\[\mathcal{G}^{i,l} = \mathcal{G}^{i,l-1} + \mathrm{LayerNorm}(A(\mathcal{Q}^{i,l}, \mathcal{K}^{i,l}, \mathcal{V}^{i,l}, \mathcal{R}^{i,l})).\]

Parameters:

rcut: float

The cut-off radius \(r_c\)

rcut_smth: float

From where the environment matrix should be smoothed \(r_s\)

sellist[int], int

list[int]: sel[i] specifies the maxmum number of type i atoms in the cut-off radius int: the total maxmum number of atoms in the cut-off radius

ntypesint

Number of element types

neuronlist[int]

Number of neurons in each hidden layers of the embedding net \(\mathcal{N}\)

axis_neuron: int

Number of the axis neuron \(M_2\) (number of columns of the sub-matrix of the embedding matrix)

tebd_dim: int

Dimension of the type embedding

tebd_input_mode: str

The input mode of the type embedding. Supported modes are [“concat”, “strip”]. - “concat”: Concatenate the type embedding with the smoothed radial information as the union input for the embedding network. - “strip”: Use a separated embedding network for the type embedding and combine the output with the radial embedding network output.

resnet_dt: bool

Time-step dt in the resnet construction: y = x + dt * phi (Wx + b)

trainable: bool

If the weights of this descriptors are trainable.

trainable_ln: bool

Whether to use trainable shift and scale weights in layer normalization.

ln_eps: float, Optional

The epsilon value for layer normalization.

type_one_side: bool

If ‘False’, type embeddings of both neighbor and central atoms are considered. If ‘True’, only type embeddings of neighbor atoms are considered. Default is ‘False’.

attn: int

Hidden dimension of the attention vectors

attn_layer: int

Number of attention layers

attn_dotr: bool

If dot the angular gate to the attention weights

attn_mask: bool

(Only support False to keep consistent with other backend references.) (Not used in this version. True option is not implemented.) If mask the diagonal of attention weights

exclude_typeslist[list[int]]

The excluded pairs of types which have no interaction with each other. For example, [[0, 1]] means no interaction between type 0 and type 1.

env_protection: float

Protection parameter to prevent division by zero errors during environment matrix calculations.

set_davg_zero: bool

Set the shift of embedding net input to zero.

activation_function: str

The activation function in the embedding net. Supported options are “gelu_tf”, “silu”, “tanh”, “gelu_erf”, “sigmoid”, “silut”, “softplus”, “relu6”, “relu”, “linear”, “none”, “gelu”.

precision: str

The precision of the embedding net parameters. Supported options are “default”, “float64”, “float32”, “bfloat16”, “float16”.

scaling_factor: float

The scaling factor of normalization in calculations of attention weights. If temperature is None, the scaling of attention weights is (N_dim * scaling_factor)**0.5

normalize: bool

Whether to normalize the hidden vectors in attention weights calculation.

temperature: float

If not None, the scaling of attention weights is temperature itself.

smooth_type_embedding: bool

Whether to use smooth process in attention weights calculation.

concat_output_tebd: bool

Whether to concat type embedding at the output of the descriptor.

stripped_type_embedding: bool, Optional

(Deprecated, kept only for compatibility.) Whether to strip the type embedding into a separate embedding network. Setting this parameter to True is equivalent to setting tebd_input_mode to ‘strip’. Setting it to False is equivalent to setting tebd_input_mode to ‘concat’. The default value is None, which means the tebd_input_mode setting will be used instead.

use_econf_tebd: bool, Optional

Whether to use electronic configuration type embedding.

use_tebd_biasbool, Optional

Whether to use bias in the type embedding layer.

type_map: list[str], Optional

A list of strings. Give the name to each type of atoms.

lmax: int

Maximum angular degree of the Cartesian moment basis. Supported values are 1 through 4.

spin

(Only support None to keep consistent with other backend references.) (Not used in this version. Not-none option is not implemented.) The old implementation of deepspin.

References

[1]

Duo Zhang, Hangrui Bi, Fu-Zhi Dai, Wanrun Jiang, Linfeng Zhang, and Han Wang. 2022. DPA-1: Pretraining of Attention-based Deep Potential Model for Molecular Simulation. arXiv preprint arXiv:2208.08236.

_update_sel_cls

_promote_degree_gain() → None

Promote the dpmodel raw degree gains from a buffer to a Parameter.

classmethod deserialize(data: dict) → DescrptDPA1

Deserialize from dict.

disable_graph_lower() → None

Persisted variant of the dpmodel escape hatch (see base class).

The buffer (and the routing bool) are PER-TASK state: multi-task share_params shares network submodules, not this buffer, so disabling the graph lower on one task branch does not propagate to branches sharing the same descriptor weights – each branch owns its routing decision.

_load_from_state_dict(state_dict: dict[str, Any], prefix: str, *args: Any, **kwargs: Any) → None

share_params(base_class: Any, shared_level: int, model_prob: float = 1.0, resume: bool = False) → None

Share parameters with base_class for multi-task training.

Level 0: share type_embedding and se_atten (all modules and buffers). Level 1: share type_embedding only.

enable_compression(min_nbor_dist: float, table_extrapolate: float = 5, table_stride_1: float = 0.01, table_stride_2: float = 0.1, check_frequency: int = -1) → None

Enable compression for the DPA1 descriptor.

Parameters:

min_nbor_dist

The nearest distance between atoms

table_extrapolate

The scale of model extrapolation

table_stride_1

The uniform stride of the first table

table_stride_2

The uniform stride of the second table

check_frequency

The overflow check frequency

_store_compress_data() → None

Store tabulated data as buffers for the compressed geometric embedding.

_store_type_embd_data() → None

Precompute type embedding outputs and store as a buffer.

call(coord_ext: torch.Tensor, atype_ext: torch.Tensor, nlist: torch.Tensor, mapping: torch.Tensor | None = None, fparam: torch.Tensor | None = None, comm_dict: dict | None = None, charge_spin: torch.Tensor | None = None) → Any

Compute the descriptor.

Parameters:

coord_ext

The extended coordinates of atoms. shape: nf x (nallx3)

atype_ext

The extended aotm types. shape: nf x nall

nlist

The neighbor list. shape: nf x nloc x nnei

mapping

The index mapping from extended to local region. not used by this descriptor.

Returns:

descriptor

The descriptor. shape: nf x nloc x (ng x axis_neuron)

gr

The rotationally equivariant and permutationally invariant single particle representation. shape: nf x nloc x ng x 3

g2

The rotationally invariant pair-partical representation. this descriptor returns None

h2

The rotationally equivariant pair-partical representation. this descriptor returns None

sw

The smooth switch function.

call_graph(graph: Any, atype: torch.Tensor, type_embedding: torch.Tensor | None = None, static_nnei: int | None = None, comm_dict: dict | None = None) → Any

Graph-native forward, routed through the fused edge kernel when eligible.

Concat and strip attention-free blocks are all graph-eligible. A geometrically compressed descriptor takes the tabulated route (_call_graph_cuda_compress() at DP_CUDA_INFER >= 1, the original fused table operator through _call_graph_compress_reference() at level 0); an uncompressed descriptor takes the embedding route (_call_graph_cuda() mega kernel, the Triton edge convolution for concat, or the dpmodel reference). static_nnei is the shape-static neighbor count the dense call() adapter supplies for the attention edge-pair enumeration; it is forwarded to the dpmodel reference and unused by the attention-free fused paths.

Unlike the dense call(), the graph lower is not @cast_precision-wrapped: edge_vec arrives in the model-agnostic fp64 .pt2 ABI. The fused CUDA kernels consume it in that leaf dtype directly (casting to the model precision in-kernel); the Triton and reference paths need the Python-side alignment below so their environment matrix runs at the model precision. Force / virial precision is decided downstream by the model output transform, not here.

supports_graph_export() → bool

Compressed DPA1 must trace its fused opaque operator.

_fused_eligible(backend: str) → bool

Whether a fused descriptor kernel can serve this block.

Parameters:

backendstr

"cuda" for the mega kernel (cuda.dpa1.graph_descriptor) or "triton" for the environment convolution (triton.dpa1). Both require the attention-free path, both tebd input modes (concat / strip) and an activation the kernel inlines (tanh / silu; see activation.ACT_CODES). cuda additionally requires no excluded types, fp32 statistics, the compiled operator library, and – since its MLP runs entirely in registers with template-specialized widths – a three-layer embedding where each layer keeps or doubles the previous width, with neuron[0] in {8, 16, 32, 64}, neuron[1] <= 64 and neuron[2] <= 128, and one inlined activation on every layer (the strip gate network is unconstrained: its table is precomputed outside the kernel). triton serves any layer stack (the head layers run on cuBLAS), so only the last activation must be inlined.

Production CUDA binaries currently instantiate only lmax=1. The lmax 2/3/4 kernel sources are retained behind the DEEPMD_ENABLE_DPA1_HIGH_LMAX CMake option and use the portable reference path unless a dedicated experimental build enables them.

_call_triton(coord_ext: torch.Tensor, atype_ext: torch.Tensor, nlist: torch.Tensor) → Any

Fused Triton environment convolution (attn-free strip / concat path).

The two embedding GEMMs stay on cuBLAS; the last layer’s activation and residual, the smooth cutoff and the moment reduction collapse into se_conv(), which never materializes an (E, ng) tensor. Strip additionally folds the type-pair gate (gathered inline); concat instead feeds the type feature through the embedding input and applies no gate. Composes under make_fx / torch.export so the operator is baked into the pt_expt .pt2.

_call_compressed(coord_ext: torch.Tensor, atype_ext: torch.Tensor, nlist: torch.Tensor) → Any

Compressed forward for DPA1 descriptor (strip only).

_call_graph_compress_reference(graph: Any, atype: torch.Tensor, type_embedding: torch.Tensor) → Any

Reference geo-compressed strip descriptor on the edge stream (attn-free).

The DP_CUDA_INFER == 0 graph path and the CPU / trace-time fallback. Evaluates the tabulated geometric embedding with the original fused table operator deepmd::tabulate_fusion_se_atten, treating each edge as a one-neighbor block (nloc = E, nnei = 1) so the operator returns the per-edge moment outer product with four or nine basis rows; a segment_sum over edge centers then forms the per-node moment. This matches the dense DescrptDPA1._call_compressed() (same table, same gate) to the fp32 summation-order floor, and composes under autograd so the force / virial assembly differentiates through it. Not make_fx-traceable (the table operator has no meta kernel); the exportable graph is the level >= 1 CUDA kernel.

Parameters:

graphNeighborGraph

Lowered neighbor graph; edge_vec is in the model precision.

atypetorch.Tensor

Flat node atom types with shape (N,), int64.

type_embeddingtorch.Tensor

Type embedding table with shape (ntypes + 1, tebd_dim).

Returns:

grrgtorch.Tensor

Descriptor with shape (N, ng * axis [+ tebd_dim]).

rot_mattorch.Tensor

Equivariant rotation matrix with shape (N, ng, 3).

_call_graph_cuda_compress(graph: Any, atype: torch.Tensor, type_embedding: torch.Tensor) → Any

Fused CUDA graph-native geo-compressed strip descriptor (attn-free).

Numerically equivalent to DescrptDPA1._call_compressed() through the dpa1_graph_compress() operator: the environment matrix, quintic table lookup, strip type-pair gate, moment reduction and G^T G contraction collapse into one CUDA mega kernel whose registered backward exposes the edge_vec gradient for the analytic force / virial assembly.

_call_graph_cuda(graph: Any, atype: torch.Tensor, type_embedding: torch.Tensor) → Any

Fused CUDA graph-native descriptor (concat, attn-free).

Numerically equivalent to DescrptDPA1DP.call_graph() through the dpa1_graph_descriptor() operator: the environment matrix, embedding MLP, moment reduction and G^T G contraction collapse into one CUDA mega kernel whose registered backward exposes the edge_vec gradient for the analytic force / virial assembly. Composes under make_fx / torch.export so the operator is baked into the pt_expt graph-form .pt2.

fused_energy_force_graph(fit: Any, graph: Any, atype: torch.Tensor, ownership: torch.Tensor, atom_bias: torch.Tensor, do_atomic_virial: bool, spin: torch.Tensor | None = None) → tuple[torch.Tensor, ...] | None

End-to-end fused energy / force / virial from the edge stream.

Collapses this descriptor, the energy fitting and the analytic force / virial assembly into one value-returning CUDA operator (no autograd tape). Returns (energy, atom_energy, force, virial, atom_virial, force_mag), with a rank-one empty magnetic-force sentinel, or None when the descriptor or fitting is not fused-eligible or the operator library is unavailable – the caller then uses the autograd lower. The geo-compressed descriptor dispatches to its tabulated operator (dpa1_graph_compress_energy_force()); the embedding-MLP descriptor to dpa1_graph_energy_force().

Parameters:

fitEnergyFittingNet

The energy fitting module fused with this descriptor.

graphNeighborGraph

The lowered neighbor graph (edge_vec, edge_index, edge_mask, n_node).

atypetorch.Tensor

Flat node atom types with shape (N,), int64.

ownershiptorch.Tensor

Energy-contributing node mask with shape (N,), bool.

atom_biastorch.Tensor

Combined fitting and atomic-model bias with shape (ntypes,).

do_atomic_virialbool

Whether to also assemble the per-atom virial.

spintorch.Tensor or None

Optional per-node magnetic moments. DPA1 does not consume them; supplying one makes the caller select a spin-capable fallback.

Returns:

tuple[torch.Tensor, …] or None

(energy, atom_energy, force, virial, atom_virial, force_mag), or None.

_call_graph_triton(graph: Any, atype: torch.Tensor, type_embedding: torch.Tensor) → Any

Fused graph-native environment convolution (concat / strip, attn-free).

Bit-exact analogue of DescrptDPA1DP.call_graph(): builds the per-edge environment matrix and embedding input, runs the two embedding GEMMs on cuBLAS, then folds the last layer, the type-pair gate (strip), the edge mask, the outer product and the segment_sum scatter into edge_conv(). The two tebd-input modes differ only in the embedding input (concat folds the type feature in; strip runs on the radial channel alone) and the gate (strip multiplies by the type-pair table, concat does not). Composes under make_fx / torch.export so the operator is baked into the pt_expt graph-form .pt2.