quantax.model.ResConv#

class quantax.model.ResConv#

Deep convolutional residual network.

__init__(nblocks: int, channels: int, kernel_size: int | ~typing.Sequence[int], sublattice: ~typing.Sequence[int] | None = None, use_rmsnorm: bool = False, final_activation: ~typing.Callable[[~jax.Array], ~numpy._typing._array_like.NDArray | ~jax.Array | ~quantax.utils.big_array.LogArray | ~quantax.utils.big_array.ScaleArray] | None = None, trans_symm: ~quantax.symmetry.symmetry.Symmetry | None = None, dtype: str | type[~typing.Any] | ~numpy.dtype | ~jax._src.typing.SupportsDType = <class 'jax.numpy.float32'>, out_dtype: str | type[~typing.Any] | ~numpy.dtype | ~jax._src.typing.SupportsDType | None = None)#

The convolutional residual network with a summation in the end.

Parameters:

  • nblocks – The number of residual blocks. Each block contains two convolutional layers.

  • channels – The number of channels. Each layer has the same amount of channels.

  • kernel_size – The kernel size. Each layer has the same kernel size.

  • sublattice – The sublattice size of the embedding, default to no sublattice.

  • use_rmsnorm – Whether to use RMSNorm in each block. Default to False, in which case a manual renormalization is applied to rescale the block input by its initial variance.

  • final_activation – The activation function in the last layer. By default, exp_by_scale is used.

  • trans_symm – The translation symmetry to be applied in the last layer, see ConvSymmetrize.

  • dtype – The data type of the parameters. Must be a real dtype.

  • out_dtype – The data type of the output wavefunction. By default, it is the same as dtype. If out_dtype is complex, pair_cpl will be applied to the output of convolutional layers to make the final output complex.

Tip

This is the recommended architecture for deep NQS.