netket.models.AbstractARNN

class netket.models.AbstractARNN(hilbert, parent=<flax.linen.module._Sentinel object>, name=None)[source]

Bases: flax.linen.module.Module

Base class for autoregressive neural networks.

Subclasses must implement __call__ and conditionals. They can also override _conditional to implement the caching for fast autoregressive sampling. See netket.nn.FastARNNConv1D for example.

Attributes
variables

Returns the variables in this module.

Return type

Mapping[str, Mapping[str, Any]]

Methods
abstract 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).