netket.nn.blocks.DeepSetMLP

class netket.nn.blocks.DeepSetMLP

Bases: Module

Implements the DeepSets architecture, which is permutation invariant and is suitable for the encoding of bosonic systems.

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

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

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 include final dimension of the network. When features_rho is None, no rho network 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__(x)

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