quantax.nn.circular_pad#

quantax.nn.circular_pad(x: Array, padding: int | Sequence[int] | Sequence[tuple[int, int]]) → Array#

Circular padding on all physical dimensions of a feature map.

The input x is given in the neighbor representation with shape (channels, Lx, Ly, ...), where the trailing axes must match the lattice shape get_lattice().shape[1:] and the leading axes are channels. Each physical dimension is padded according to padding, wrapping around the boundary for periodic or anti-periodic boundary conditions.

The wrap obeys the lattice periodicity in the original representation, not a naive per-axis wrap of the neighbor grid: for skewed arrangements such as TriangularB a plain wrap is only covariant under a subgroup of the translations. The correct (possibly twisted) wrap is derived automatically from to_neighbor_repr, assuming one step on the neighbor grid translates all sites by the same lattice vector. It reduces to a plain wrap when the neighbor grid is an unskewed torus (e.g. square lattices, whose to_neighbor_repr is the identity).

Implemented as a single gather with a statically-built index, which is memory-bandwidth optimal and noticeably faster than jnp.pad(mode="wrap").

Parameters:

  • x – The feature map of shape (channels, Lx, Ly, ...) in neighbor representation, with spatial shape equal to get_lattice().shape[1:].

  • padding –

    The padding on every physical dimension, following the convention of equinox.nn.Conv (string inputs are not supported). It can be

    • an int: the same padding on both sides of every dimension;

    • a sequence of int: one symmetric padding per dimension;

    • a sequence of (low, high) pairs: asymmetric padding per dimension.