quantax.nn.Conv#

class quantax.nn.Conv#

Bases: Module

Convolution layer whose padding is determined by the lattice boundary.

This is analogous to equinox.nn.Conv, but it has no padding or padding_mode arguments: the padding is fixed by the boundary conditions of the current lattice (get_lattice), and the number of spatial dimensions is taken from lattice.ndim. For each physical dimension,

  • periodic and anti-periodic dimensions (boundary != 0) are wrapped with circular_pad (a “SAME”-sized wrap, no anti-periodic sign);

  • open dimensions (boundary == 0) are zero-padded (“SAME”) inside jax.lax.conv_with_general_padding.

The output therefore has the same spatial shape as the input (for unit stride). The input/output are given in the neighbor representation.

__init__(in_channels: int, out_channels: int, kernel_size: int | ~typing.Sequence[int], stride: int | ~typing.Sequence[int] = 1, dilation: int | ~typing.Sequence[int] = 1, use_bias: bool = True, dtype: str | type[~typing.Any] | ~numpy.dtype | ~jax._src.typing.SupportsDType = <class 'jax.numpy.float32'>, *, key: ~jaxtyping.Key)#
Parameters:

  • in_channels – The number of input channels.

  • out_channels – The number of output channels.

  • kernel_size – The size of the convolutional kernel, an int or one int per physical dimension.

  • stride – The stride of the convolution, by default 1.

  • dilation – The dilation of the convolution, by default 1.

  • use_bias – Whether to add a bias after the convolution, by default True.

  • dtype – The data type of the parameters, by default float32.

  • key – The PRNG key for weight initialization.

__call__(x: Array) → Array#
Parameters:

x – The input of shape (in_channels, *spatial) in neighbor representation.