deepmd.pt_expt.utils.network

Classes

TorchArrayParam	Parameter subclass that supports np.array(param) conversion.
NativeLayer	PyTorch layer wrapping dpmodel's NativeLayer.
NativeNet	Native representation of a neural network.
EmbeddingNet	The embedding network.
FittingNet	The fitting network. It may be implemented as an embedding
NetworkCollection	A collection of networks for multiple elements.
LayerNorm	Implementation of Layer Normalization layer.
Identity	The unit operation of a native model.

Functions

_torch_activation(→ torch.Tensor)

Apply activation function using native torch ops.

Module Contents

class deepmd.pt_expt.utils.network.TorchArrayParam

Bases: torch.nn.Parameter

Parameter subclass that supports np.array(param) conversion.

Note: this class is intentionally NOT used for model parameters. make_fx (torch.fx.experimental.proxy_tensor) uses ProxyTorchDispatchMode to intercept tensor operations. When an operand is a subclass of torch.Tensor (including subclasses of torch.nn.Parameter), PyTorch invokes the __torch_function__ protocol which the proxy dispatch mode does not handle, causing aten.mm and other ops to fail with “Multiple dispatch failed … returned NotImplemented”. Using plain torch.nn.Parameter avoids this because the proxy mode is designed to work with the base Parameter type. TorchArrayParam is kept only for backward compatibility and should not be used for new code.

__array__(dtype: Any | None = None) → numpy.ndarray

class deepmd.pt_expt.utils.network.NativeLayer(*args: Any, **kwargs: Any)

Bases: deepmd.dpmodel.utils.network.NativeLayer, torch.nn.Module

PyTorch layer wrapping dpmodel’s NativeLayer.

Two aspects of the inherited dpmodel call() are incompatible with make_fx tracing (used to export forward_lower with autograd.grad-based force/virial computation):

1. Ellipsis indexing (self.w[...]): On a torch.Tensor this triggers aten.alias, an op that ProxyTorchDispatchMode does not support, resulting in “Multiple dispatch failed for aten.alias.default”.

2. ``array_api_compat`` wrappers (xp = array_api_compat .array_namespace(x); xp.matmul(…)): The wrappers re-enter torch.matmul through Python, which goes through the __torch_function__ protocol. Under the proxy dispatch mode this path also fails with “Multiple dispatch failed”.

This class therefore overrides call() with an implementation that uses plain torch ops exclusively (torch.matmul, torch.tanh, etc.), avoiding both issues.

Trainable weights are stored as plain torch.nn.Parameter (not TorchArrayParam) for the same make_fx compatibility reason — see the TorchArrayParam docstring.

__call__(*args: Any, **kwargs: Any) → Any

Forward pass in NumPy implementation.

__setattr__(name: str, value: Any) → None

call(x: torch.Tensor) → torch.Tensor

Forward pass using pure torch ops.

Overrides dpmodel’s call() to ensure compatibility with make_fx (torch.fx.experimental.proxy_tensor).

The dpmodel implementation uses self.w[...] and array_api_compat.array_namespace(x).matmul(…) for backend-agnostic array operations. Both patterns break under make_fx’s ProxyTorchDispatchMode:

• self.w[...] emits aten.alias which the proxy mode cannot dispatch.

• array_api_compat re-enters torch.matmul via Python, hitting __torch_function__ which the proxy mode returns NotImplemented for.

This override uses torch.matmul, torch.cat, and _torch_activation directly, sidestepping both issues.

forward(x: torch.Tensor) → torch.Tensor

deepmd.pt_expt.utils.network._torch_activation(x: torch.Tensor, name: str) → torch.Tensor

Apply activation function using native torch ops.

The dpmodel get_activation_fn returns closures that call array_api_compat.array_namespace(x).tanh(x) etc. Under make_fx proxy tracing, the array_api_compat indirection triggers __torch_function__ dispatch failures. This function calls torch.tanh and friends directly to avoid the issue.

class deepmd.pt_expt.utils.network.NativeNet(layers: list[dict] | None = None)

Bases: make_multilayer_network(NativeLayer, NativeOP)

Native representation of a neural network.

Parameters:

layerslist[NativeLayer], optional

The layers of the network.

layers

class deepmd.pt_expt.utils.network.EmbeddingNet(*args: Any, **kwargs: Any)

Bases: deepmd.dpmodel.utils.network.EmbeddingNet, torch.nn.Module

The embedding network.

Parameters:

in_dim

Input dimension.

neuron

The number of neurons in each layer. The output dimension is the same as the dimension of the last layer.

activation_function

The activation function.

resnet_dt

Use time step at the resnet architecture.

precision

Floating point precision for the model parameters.

seedint, optional

Random seed.

biasbool, Optional

Whether to use bias in the embedding layer.

trainablebool or list[bool], Optional

Whether the weights are trainable. If a list, each element corresponds to a layer.

layers

__call__(*args: Any, **kwargs: Any) → Any

forward(x: torch.Tensor) → torch.Tensor

class deepmd.pt_expt.utils.network.FittingNet(*args: Any, **kwargs: Any)

Bases: deepmd.dpmodel.utils.network.FittingNet, torch.nn.Module

The fitting network. It may be implemented as an embedding net connected with a linear output layer.

Parameters:

in_dim

Input dimension.

out_dim

Output dimension

neuron

The number of neurons in each hidden layer.

activation_function

The activation function.

resnet_dt

Use time step at the resnet architecture.

precision

Floating point precision for the model parameters.

bias_out

The last linear layer has bias.

seedint, optional

Random seed.

trainablebool or list[bool], optional

Whether the network is trainable.

layers

__call__(*args: Any, **kwargs: Any) → Any

forward(x: torch.Tensor) → torch.Tensor

class deepmd.pt_expt.utils.network.NetworkCollection(*args: Any, **kwargs: Any)

Bases: deepmd.dpmodel.utils.network.NetworkCollection

A collection of networks for multiple elements.

The number of dimensions for types might be 0, 1, or 2. - 0: embedding or fitting with type embedding, in () - 1: embedding with type_one_side, or fitting, in (type_i) - 2: embedding without type_one_side, in (type_i, type_j)

Parameters:

ndimint

The number of dimensions.

network_typestr, optional

The type of the network.

networksdict, optional

The networks to initialize with.

NETWORK_TYPE_MAP: ClassVar[dict[str, type]]

_module_networks

__setitem__(key: int | tuple, value: Any) → None

class deepmd.pt_expt.utils.network.LayerNorm(num_in: int, eps: float = 1e-05, uni_init: bool = True, trainable: bool = True, precision: str = DEFAULT_PRECISION, seed: int | list[int] | None = None)

Bases: deepmd.dpmodel.utils.network.LayerNorm, NativeLayer

Implementation of Layer Normalization layer.

Parameters:

num_inint

The input dimension of the layer.

epsfloat, optional

A small value added to prevent division by zero in calculations.

uni_initbool, optional

If initialize the weights to be zeros and ones.

trainablebool, optional

If the weights are trainable.

precisionstr, optional

The precision of the layer.

seedint, optional

Random seed.

class deepmd.pt_expt.utils.network.Identity

Bases: deepmd.dpmodel.utils.network.Identity

The unit operation of a native model.

forward(x: torch.Tensor) → torch.Tensor