deepmd.jax.descriptor.repformers

deepmd.jax.descriptor.repformers#

Classes#

`DescrptBlockRepformers`	The repformer descriptor block.
`Atten2Map`	The unit operation of a native model.
`Atten2MultiHeadApply`	The unit operation of a native model.
`Atten2EquiVarApply`	The unit operation of a native model.
`LocalAtten`	The unit operation of a native model.
`RepformerLayer`	The unit operation of a native model.

Module Contents#

class deepmd.jax.descriptor.repformers.DescrptBlockRepformers(rcut, rcut_smth, sel: int, ntypes: int, nlayers: int = 3, g1_dim=128, g2_dim=16, axis_neuron: int = 4, direct_dist: bool = False, update_g1_has_conv: bool = True, update_g1_has_drrd: bool = True, update_g1_has_grrg: bool = True, update_g1_has_attn: bool = True, update_g2_has_g1g1: bool = True, update_g2_has_attn: bool = True, update_h2: bool = False, attn1_hidden: int = 64, attn1_nhead: int = 4, attn2_hidden: int = 16, attn2_nhead: int = 4, attn2_has_gate: bool = False, activation_function: str = 'tanh', update_style: str = 'res_avg', update_residual: float = 0.001, update_residual_init: str = 'norm', set_davg_zero: bool = True, smooth: bool = True, exclude_types: list[tuple[int, int]] = [], env_protection: float = 0.0, precision: str = 'float64', trainable_ln: bool = True, use_sqrt_nnei: bool = True, g1_out_conv: bool = True, g1_out_mlp: bool = True, ln_eps: float | None = 1e-05, seed: int | list[int] | None = None)[source]#

Bases: deepmd.dpmodel.descriptor.repformers.DescrptBlockRepformers

The repformer descriptor block.

Parameters:

rcutfloat: The cut-off radius.
rcut_smthfloat: Where to start smoothing. For example the 1/r term is smoothed from rcut to rcut_smth.
selint: Maximally possible number of selected neighbors.
ntypesint: Number of element types
nlayersint, optional: Number of repformer layers.
g1_dimint, optional: Dimension of the first graph convolution layer.
g2_dimint, optional: Dimension of the second graph convolution layer.
axis_neuronint, optional: Size of the submatrix of G (embedding matrix).
direct_distbool, optional: Whether to use direct distance information (1/r term) in the repformer block.
update_g1_has_convbool, optional: Whether to update the g1 rep with convolution term.
update_g1_has_drrdbool, optional: Whether to update the g1 rep with the drrd term.
update_g1_has_grrgbool, optional: Whether to update the g1 rep with the grrg term.
update_g1_has_attnbool, optional: Whether to update the g1 rep with the localized self-attention.
update_g2_has_g1g1bool, optional: Whether to update the g2 rep with the g1xg1 term.
update_g2_has_attnbool, optional: Whether to update the g2 rep with the gated self-attention.
update_h2bool, optional: Whether to update the h2 rep.
attn1_hiddenint, optional: The hidden dimension of localized self-attention to update the g1 rep.
attn1_nheadint, optional: The number of heads in localized self-attention to update the g1 rep.
attn2_hiddenint, optional: The hidden dimension of gated self-attention to update the g2 rep.
attn2_nheadint, optional: The number of heads in gated self-attention to update the g2 rep.
attn2_has_gatebool, optional: Whether to use gate in the gated self-attention to update the g2 rep.
activation_functionstr, optional: The activation function in the embedding net.
update_stylestr, optional: Style to update a representation. Supported options are: -‘res_avg’: Updates a rep u with: u = 1/\sqrt{n+1} (u + u_1 + u_2 + … + u_n) -‘res_incr’: Updates a rep u with: u = u + 1/\sqrt{n} (u_1 + u_2 + … + u_n) -‘res_residual’: Updates a rep u with: u = u + (r1*u_1 + r2*u_2 + … + r3*u_n) where r1, r2 … r3 are residual weights defined by update_residual and update_residual_init.
update_residualfloat, optional: When update using residual mode, the initial std of residual vector weights.
update_residual_initstr, optional: When update using residual mode, the initialization mode of residual vector weights.
set_davg_zerobool, optional: Set the normalization average to zero.
precisionstr, optional: The precision of the embedding net parameters.
smoothbool, optional: Whether to use smoothness in processes such as attention weights calculation.
exclude_typeslist[list[int]], optional: 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_protectionfloat, optional: Protection parameter to prevent division by zero errors during environment matrix calculations. For example, when using paddings, there may be zero distances of neighbors, which may make division by zero error during environment matrix calculations without protection.
trainable_lnbool, optional: Whether to use trainable shift and scale weights in layer normalization.
use_sqrt_nneibool, optional: Whether to use the square root of the number of neighbors for symmetrization_op normalization instead of using the number of neighbors directly.
g1_out_convbool, optional: Whether to put the convolutional update of g1 separately outside the concatenated MLP update.
g1_out_mlpbool, optional: Whether to put the self MLP update of g1 separately outside the concatenated MLP update.
ln_epsfloat, optional: The epsilon value for layer normalization.
seedint, optional: The random seed for initialization.

__setattr__(name: str, value: Any) → None[source]#

class deepmd.jax.descriptor.repformers.Atten2Map(input_dim: int, hidden_dim: int, head_num: int, has_gate: bool = False, smooth: bool = True, attnw_shift: float = 20.0, precision: str = 'float64', seed: int | list[int] | None = None)[source]#

Bases: deepmd.dpmodel.descriptor.repformers.Atten2Map