deepmd.jax.utils.network

deepmd.jax.utils.network#

Classes#

`ArrayAPIParam`
`NativeLayer`	Native representation of a layer.
`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.

Module Contents#

class deepmd.jax.utils.network.ArrayAPIParam[source]#

Bases: deepmd.jax.env.nnx.Param

__array__(*args, **kwargs)[source]#

__array_namespace__(*args, **kwargs)[source]#

__dlpack__(*args, **kwargs)[source]#

__dlpack_device__(*args, **kwargs)[source]#

class deepmd.jax.utils.network.NativeLayer(num_in, num_out, bias: bool = True, use_timestep: bool = False, activation_function: str | None = None, resnet: bool = False, precision: str = DEFAULT_PRECISION, seed: int | list[int] | None = None)[source]#

Bases: deepmd.dpmodel.utils.network.NativeLayer

Native representation of a layer.

Parameters:

wnp.ndarray, optional: The weights of the layer.
bnp.ndarray, optional: The biases of the layer.
idtnp.ndarray, optional: The identity matrix of the layer.
activation_functionstr, optional: The activation function of the layer.
resnetbool, optional: Whether the layer is a residual layer.
precisionstr, optional: The precision of the layer.
seedint, optional: Random seed.

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

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

Bases: make_multilayer_network(NativeLayer, NativeOP)

Native representation of a neural network.

Parameters:

layerslist[NativeLayer], optional: The layers of the network.

class deepmd.jax.utils.network.EmbeddingNet(in_dim, neuron: list[int] = [24, 48, 96], activation_function: str = 'tanh', resnet_dt: bool = False, precision: str = DEFAULT_PRECISION, seed: int | list[int] | None = None, bias: bool = True)[source]#

Bases: make_embedding_network(NativeNet, NativeLayer)

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.

class deepmd.jax.utils.network.FittingNet(in_dim, out_dim, neuron: list[int] = [24, 48, 96], activation_function: str = 'tanh', resnet_dt: bool = False, precision: str = DEFAULT_PRECISION, bias_out: bool = True, seed: int | list[int] | None = None)[source]#

Bases: make_fitting_network(EmbeddingNet, NativeNet, NativeLayer)

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.

class deepmd.jax.utils.network.NetworkCollection(ndim: int, ntypes: int, network_type: str = 'network', networks: list[NativeNet | dict] = [])[source]#

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]][source]#

class deepmd.jax.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)[source]#

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.