quantax.nn.ReshapeConv

class quantax.nn.ReshapeConv

Bases: Module

Reshape the input to the shape suitable for convolutional layers.

A Fock state in Quantax is usually given by a 1D array with entries +1/-1. This layer reshapes it to shape.

__init__(dtype: str | type[Any] | dtype | SupportsDType = <class 'jax.numpy.float32'>)

Parameters:

dtype – Convert the input to the given data type, by default float32.

__call__(x: Array) → Array

Reshape a configuration to the lattice shape and cast it to dtype.