netket.models.ARNNConv2D#

class netket.models.ARNNConv2D[source]#

Bases: netket.models.autoreg.AbstractARNN

Autoregressive neural network with 2D convolution layers.

Attributes
kernel_dilation: Tuple[int, int] = (1, 1)#

a sequence of 2 integers, giving the dilation factor to apply in each spatial dimension of the convolution kernel (default: 1).

machine_pow: int = 2#

exponent to normalize the outputs of __call__.

precision: Any = None#

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

use_bias: bool = True#

True).

Type

whether to add a bias to the output (default

variables#

Returns the variables in this module.

Return type

Mapping[str, Mapping[str, Any]]

layers: int#

number of layers.

features: Union[Iterable[int], int]#

number of features in each layer. If a single number is given, all layers except the last one will have the same number of features.

kernel_size: Tuple[int, int]#

shape of the convolutional kernel (h, w). Typically, h = w // 2 + 1.

hilbert: netket.hilbert.HomogeneousHilbert#

the Hilbert space. Only homogeneous unconstrained Hilbert spaces are supported.

Methods
activation()#

Scaled exponential linear unit activation.

Computes the element-wise function:

\[\begin{split}\mathrm{selu}(x) = \lambda \begin{cases} x, & x > 0\\ \alpha e^x - \alpha, & x \le 0 \end{cases}\end{split}\]

where \(\lambda = 1.0507009873554804934193349852946\) and \(\alpha = 1.6732632423543772848170429916717\).

For more information, see Self-Normalizing Neural Networks.

Parameters

x (Any) – input array

Return type

Any

bias_init(shape, dtype=<class 'jax.numpy.float64'>)#

An initializer that returns a constant array full of zeros.

The key argument is ignored.

>>> import jax, jax.numpy as jnp
>>> jax.nn.initializers.zeros(jax.random.PRNGKey(42), (2, 3), jnp.float32)
DeviceArray([[0., 0., 0.],
             [0., 0., 0.]], dtype=float32)
Parameters

dtype (Any) –

conditionals(inputs)[source]#

Computes the conditional probabilities for each site to take each value.

Parameters

inputs (Union[ndarray, DeviceArray, Tracer]) – configurations with dimensions (batch, Hilbert.size).

Return type

Union[ndarray, DeviceArray, Tracer]

Returns

The probabilities with dimensions (batch, Hilbert.size, Hilbert.local_size).

Examples

>>> import pytest; pytest.skip("skip automated test of this docstring")
>>>
>>> p = model.apply(variables, σ, method=model.conditionals)
>>> print(p[2, 3, :])
[0.3 0.7]
# For the 3rd spin of the 2nd sample in the batch,
# it takes probability 0.3 to be spin down (local state index 0),
# and probability 0.7 to be spin up (local state index 1).
has_rng(name)#

Returns true if a PRNGSequence with name name exists.

Return type

bool

Parameters

name (str) –

kernel_init(shape, dtype=<class 'jax.numpy.float64'>)#
put_variable(col, name, value)#

Sets the value of a Variable.

Parameters
  • col (str) – the variable collection.

  • name (str) – the name of the variable.

  • value (Any) – the new value of the variable.

Returns: