Potential neural networks (stories.potentials)

class stories.potentials.MLPPotential(features: ~typing.Sequence[int] = (128, 128), activation: callable = <function gelu>, parent: ~typing.Type[~flax.linen.module.Module] | ~flax.core.scope.Scope | ~typing.Type[~flax.linen.module._Sentinel] | None = <flax.linen.module._Sentinel object>, name: str | None = None)

Bases: Module

This class defines a simple multi-layer perceptron (MLP) potential which takes a batch of cells and returns a batch of scalars. The activation function is a gelu function, which is a smooth approximation to the rectified linear unit (ReLU). This makes the gradient of the potential twice differentiable wrt input.

Parameters:

  • features – A sequence of integers specifying the number of hidden units in each

  • hidden (layer of the MLP. The length of this sequence determines the number of)

  • MLP. (layers in the)

  • activation – Activation function to use in the hidden layers.

activation(approximate: bool = True) → Array

Gaussian error linear unit activation function.

If approximate=False, computes the element-wise function:

\[\mathrm{gelu}(x) = \frac{x}{2} \left(1 + \mathrm{erf} \left( \frac{x}{\sqrt{2}} \right) \right)\]

If approximate=True, uses the approximate formulation of GELU:

\[\mathrm{gelu}(x) = \frac{x}{2} \left(1 + \mathrm{tanh} \left( \sqrt{\frac{2}{\pi}} \left(x + 0.044715 x^3 \right) \right) \right)\]

For more information, see Gaussian Error Linear Units (GELUs), section 2.

Parameters:

  • x – input array

  • approximate – whether to use the approximate or exact formulation.

features: Sequence[int] = (128, 128)
name: str | None = None
parent: Type[Module] | Scope | Type[_Sentinel] | None = None
scope: Scope | None = None