netket.nn.FastMaskedConv2D

class netket.nn.FastMaskedConv2D

Bases: Module

2D convolution module with mask for fast autoregressive NN.

See netket.models.FastARNNSequential for a brief explanation of fast autoregressive sampling.

Attributes

feature_group_count: int = 1

1).

Type:

if specified, divides the input features into groups (default

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

L: int

edge length of the 2D lattice.

features: int

number of convolution filters.

kernel_size: tuple[int, int]

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

kernel_dilation: tuple[int, int]

a sequence of 2 integers, giving the dilation factor to apply in each spatial dimension of the convolution kernel.

exclusive: bool

True if an output element does not depend on the input element at the same index.

kernel_init: Callable[[Any, Sequence[int], None | str | type[Any] | dtype | _SupportsDType], Array]

initializer for the convolutional kernel.

bias_init: Callable[[Any, Sequence[int], None | str | type[Any] | dtype | _SupportsDType], Array]

initializer for the bias.

Methods

__call__(inputs)

Applies the masked convolution to all input sites.

Parameters:

inputs (Union[ndarray, Array]) – input data with dimensions (batch, width, height, features).

Return type:

Union[ndarray, Array]

Returns:

The convolved data.

update_site(inputs, index)

Adds an input site into the cache, and applies the masked convolution to the cache.

Parameters:

• inputs (Union[ndarray, Array]) – an input site to be added into the cache with dimensions (batch, features).

• index (int) – the index of the output site. The index of the input site should be index - self.exclusive.

Return type:

Union[ndarray, Array]

Returns:

The next output site with dimensions (batch, features).