deepmd.dpmodel.descriptor.dpa4_nn.so3

deepmd.dpmodel.descriptor.dpa4_nn.so3#

SO(3)-equivariant linear layers for DPA4/SeZM.

This module is the dpmodel port of deepmd.pt.model.descriptor.sezm_nn.so3. It defines the channel-only and focus-aware linear maps used by the DPA4 SO(3) feature transformations. All three pt classes are ported: FocusLinear (used by so2, grid_net, activation), ChannelLinear (used by so2, grid_net), and SO3Linear (used by so2, ffn).

Serialization contract: SO3Linear mirrors the pt serialize() format exactly (same config and @variables keys), so pt serialize() output deserializes directly. The pt FocusLinear and ChannelLinear define no serialize() (they only appear nested inside larger modules’ state_dicts); their dpmodel serialize()/deserialize() use @variables keys equal to the pt state_dict key names (weight, bias) so that pt state-dict fragments load directly.

Weight initialization is distribution-equivalent to the pt version (drawn from np.random.default_rng instead of the torch generator stream), the same convention as utils.init_trunc_normal_fan_in_out.

Classes#

`FocusLinear`	Per-focus linear projection on the last feature axis.
`ChannelLinear`	Channel-only linear projection on the last feature axis.
`SO3Linear`	Focus-aware degree-wise linear self-interaction.

Module Contents#

class deepmd.dpmodel.descriptor.dpa4_nn.so3.FocusLinear(*, in_channels: int, out_channels: int, n_focus: int, precision: str = DEFAULT_PRECISION, bias: bool = True, trainable: bool = True, seed: int | list[int] | None = None, init_std: float | None = None)[source]#

Bases: deepmd.dpmodel.NativeOP

Per-focus linear projection on the last feature axis.

Parameters:

in_channelsint: Input feature dimension.
out_channelsint: Output feature dimension.
n_focusint: Number of focus streams.
precisionstr: Parameter precision.
biasbool: Whether to use bias.
trainablebool: Whether parameters are trainable.
seedint | list[int] | None: Random seed for initialization.
init_stdfloat | None: If given, use normal(0, init_std) instead of default uniform init. Useful for gate projections where small initial logits are desired.

Notes

Parameters are stored in (in, out) convention to match Muon’s rectangular correction assumption (rows=fan_in, cols=fan_out): - weight: (in_channels, n_focus * out_channels) - bias: (n_focus * out_channels,)

in_channels[source]#

out_channels[source]#

n_focus[source]#

precision = 'float64'[source]#

trainable = True[source]#

use_bias = True[source]#

weight[source]#

call(x: Any) → Any[source]#

Apply the per-focus linear projection.

Parameters:

xArray: Input array with shape (B, F, Cin).

Returns:

Array: Projected array with shape (B, F, Cout).

serialize() → dict[str, Any][source]#

Serialize the FocusLinear to a dict.

The pt FocusLinear has no serialize(); the @variables keys here match the pt state_dict key names (weight, bias).

classmethod deserialize(data: dict[str, Any]) → FocusLinear[source]#: Deserialize a FocusLinear from a dict.

class deepmd.dpmodel.descriptor.dpa4_nn.so3.ChannelLinear(*, in_channels: int, out_channels: int, precision: str = DEFAULT_PRECISION, bias: bool = True, trainable: bool = True, seed: int | list[int] | None = None, init_std: float | None = None)[source]#

Bases: deepmd.dpmodel.NativeOP

Channel-only linear projection on the last feature axis.

Parameters:

in_channelsint: Input feature dimension.
out_channelsint: Output feature dimension.
precisionstr: Parameter precision.
biasbool: Whether to use bias.
trainablebool: Whether parameters are trainable.
seedint | list[int] | None: Random seed for initialization.
init_stdfloat | None: If given, use normal(0, init_std) instead of default uniform init. Useful for gate projections where small initial logits are desired.

Notes

Parameters are stored in (in, out) convention to match Muon’s rectangular correction assumption (rows=fan_in, cols=fan_out): - weight: (in_channels, out_channels) - bias: (out_channels,)

in_channels[source]#

out_channels[source]#

precision = 'float64'[source]#

trainable = True[source]#

use_bias = True[source]#

weight[source]#

call(x: Any) → Any[source]#

Apply the channel-only linear projection.

Parameters:

xArray: Input array with shape (..., C_in).

Returns:

Array: Projected array with shape (..., C_out).

serialize() → dict[str, Any][source]#

Serialize the ChannelLinear to a dict.

The pt ChannelLinear has no serialize(); the @variables keys here match the pt state_dict key names (weight, bias).

classmethod deserialize(data: dict[str, Any]) → ChannelLinear[source]#: Deserialize a ChannelLinear from a dict.

class deepmd.dpmodel.descriptor.dpa4_nn.so3.SO3Linear(*, lmax: int, in_channels: int, out_channels: int, n_focus: int = 1, precision: str = DEFAULT_PRECISION, mlp_bias: bool = False, trainable: bool = True, seed: int | list[int] | None = None, init_std: float | None = None)[source]#

Bases: deepmd.dpmodel.NativeOP

Focus-aware degree-wise linear self-interaction.

The key insight is that weights are shared across all m components within each l block.

Parameters:

lmaxint: Maximum spherical harmonic degree.
in_channelsint: Number of input channels per (l, m) coefficient.
out_channelsint: Number of output channels per (l, m) coefficient.
n_focusint: Number of focus streams.
precisionstr: Parameter precision.
mlp_biasbool: Whether to use bias for l=0 (scalar) components.
trainablebool: Whether parameters are trainable.
seedint | list[int] | None: Random seed for weight initialization.
init_stdfloat | None: If given, use normal(0, init_std) for all weights instead of default trunc-normal fan-in/fan-out init. Use 0.0 for zero initialization.

Notes

Weight storage: (lmax+1, C_in, F*C_out).
Bias storage: (F*C_out,), only applied to l=0 scalar components.
Runtime view restores weights to (lmax+1, C_in, F, C_out) via reshape.
expand_index maps each packed (l,m) position to its l value.
The pt einsum ndfi,difo->ndfo is expressed as a broadcast batched matmul, which keeps the whole multi-focus path vectorized.

lmax[source]#

in_channels[source]#

out_channels[source]#

n_focus = 1[source]#

precision = 'float64'[source]#

trainable = True[source]#

ebed_dim = 1[source]#

mlp_bias = False[source]#

weight[source]#

expand_index[source]#

call(x: Any) → Any[source]#

Apply the degree-wise linear self-interaction.

Parameters:

xArray: Input features with shape (N, D, F, C_in) where D=(lmax+1)^2.

Returns:

Array: Order-wise mixed features with shape (N, D, F, C_out).

serialize() → dict[str, Any][source]#: Serialize the SO3Linear to a dict (pt-compatible format).

classmethod deserialize(data: dict[str, Any]) → SO3Linear[source]#: Deserialize an SO3Linear from a dict.