netket.nn#

This sub-module extends flax.linen with layers and tools that are useful to applications in quantum physics. Read more about the design goal of this module in their README

Linear Modules#

DenseSymm

Implements a projection onto a symmetry group.

DenseEquivariant

A group convolution operation that is equivariant over a symmetry group.

MaskedDense1D

1D linear transformation module with mask for autoregressive NN.

MaskedConv1D

1D convolution module with mask for autoregressive NN.

MaskedConv2D

2D convolution module with mask for autoregressive NN.

FastMaskedDense1D

1D linear transformation module with mask for fast autoregressive NN.

FastMaskedConv1D

1D convolution module with mask for fast autoregressive NN.

FastMaskedConv2D

2D convolution module with mask for fast autoregressive NN.

Activation functions#

activation.reim(f)

Modifies a non-linearity to act separately on the real and imaginary parts

activation.reim_relu

Returns the relu non-linearity, applied separately to the real and imaginary parts

activation.reim_selu()

selu applied separately to the real andimaginary parts of it's input.

activation.log_cosh(x)

Logarithm of the hyperbolic cosine, implemented in a more stable way.

activation.log_sinh(x)

Logarithm of the hyperbolic sine.

activation.log_tanh(x)

Logarithm of the hyperbolic tangent.

Miscellaneous Functions#

binary_encoding(hilbert, x, *[, max_bits])

Encodes the array x into a set of binary-encoded variables described by the shape of a Hilbert space.

states_to_numbers(hilbert, σ)

Converts the configuration σ to a 64-bit integer denoting its index in the full Hilbert space.

Utility functions#

to_array(hilbert, apply_fun, variables[, ...])

Computes apply_fun(variables, states) on all states of hilbert and returns

to_matrix(hilbert, machine, params[, normalize])

Blocks#

blocks.MLP

A Multi-Layer Perceptron with hidden layers.

blocks.DeepSetMLP

Implements the DeepSets architecture, which is permutation invariant.

blocks.SymmExpSum

A flax module symmetrizing the log-wavefunction \(\log\psi_\theta(\sigma)\) encoded into another flax module (flax.linen.Module) by summing over all possible symmetries \(g\) in a certain discrete permutation group \(G\).