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: 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: 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: Callable | None = None

The nonlinear activation function at the output layer.

precision: Precision | None = 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: Callable | None

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], None | str | type[Any] | dtype | _SupportsDType], Array]

Initializer for the Dense layer matrix

bias_init: Callable[[Any, Sequence[int], None | str | type[Any] | dtype | _SupportsDType], 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.