quantax.nn.Embedding#

class quantax.nn.Embedding#

Bases: Module

Learned positional encoding of a configuration onto the lattice grid.

The positional encoding PE holds one embedding table per position in the sublattice, tiled over the spatial grid lattice.shape[1:]. Each unit cell’s state (see input_to_index) selects an entry from the table at its position, so the output is fully indexed from PE.

__init__(d: int, sublattice: Sequence[int] | None = None, dtype: str | type[Any] | dtype | SupportsDType = <class 'jax.numpy.float32'>)#
Parameters:

  • d – The embedding dimension (number of output channels per site).

  • sublattice – The spatial period of the positional encoding PE, tiled over the grid. Default to None, i.e. the same encoding on all positions.

  • dtype – The data type of the parameters. Default to float32.

__call__(s: Array) → Array#

Embed a configuration into an array of shape (d, *lattice.shape[1:]).