netket.models.DeepSetMLP

class netket.models.DeepSetMLP

Bases: Module

Implements the DeepSets architecture, which is permutation invariant.

\[f(x_1,...,x_N) = \rho\left(\sum_i \phi(x_i)\right)\]

that is suitable for the simulation of bosonic.

The input shape must have an axis that is reshaped to (…, N, D), where we pool over N.

See DeepSetRelDistance for the bosonic wave function ansatz in https://arxiv.org/abs/1703.06114

Attributes

features_phi: Union[int, tuple[int, ...], None] = None

Number of features in each layer for phi network. When features_phi is None, no phi network is created.

features_rho: Union[int, tuple[int, ...], None] = None

Number of features in each layer for rho network. Should not include the final layer of dimension 1, which is included automatically. When features_rho is None, a single layer MLP with output 1 is created.

output_activation: Optional[Callable] = None

The nonlinear activation function at the output layer.

precision: Optional[Precision] = None

numerical precision of the computation see jax.lax.Precision for details.

use_bias: bool = True

if True uses a bias in all layers.

hidden_activation: Optional[Callable]

The nonlinear activation function between hidden layers.

pooling: Callable

The pooling operation to be used after the phi-transformation

kernel_init: Callable[[Any, Sequence[int], Any], Union[ndarray, Array]]

Initializer for the Dense layer matrix

bias_init: Callable[[Any, Sequence[int], Any], Union[ndarray, Array]]

Initializer for the hidden bias

Methods

__call__(input)

Call self as a function.