nn#

Modules#

`Sequential`(layers[, holomorphic])	A sequence of `equinox.Module` applied in order similar to Sequential in Equinox.
`RawInputLayer`()	The layer that takes not only the output of the previous layer, but also the raw input basis state.
`RefModel`()	The model that allows accelerated forward pass through local updates and internal quantities.

`sinhp1_by_scale`(x)	\(f(x) = \sinh(x) + 1\).
`prod_by_log`(x)	\(f(x) = \prod x\).
`exp_by_scale`(x)	\(f(x) = \exp(x)\).
`exp_by_log`(x)	\(f(x) = \exp(x)\).
`crelu`(x)	Complex relu activation function \(f(x) = \mathrm{ReLU}(\mathrm{Re}\,x) + i\,\mathrm{ReLU}(\mathrm{Im}\,x)\).
`cardioid`(x)	\(f(z) = \frac{1}{2}(1 + \cos\phi)\,z\), where \(\phi = \arg z\).
`pair_cpl`(x)	\(f(x) = x_1 + i x_2\) , where \(x = (x_1, x_2)\).

The functions below take an equinox.nn.Linear or equinox.nn.Conv layer, initialize its weight with the given scheme and set its bias to 0.

`apply_lecun_normal`(-> ~equinox.nn._linear.Linear)	Apply the Lecun normal initializer.
`apply_he_normal`(-> ~equinox.nn._linear.Linear)	Apply the He normal initializer.
`apply_glorot_normal`(...)	Apply the Glorot normal initializer.

Quantax also re-exports the JAX initializers with their axes fixed to the convention of equinox.nn.Linear and equinox.nn.Conv weights (in_axis=1, out_axis=0, batch_axis=()), namely variance_scaling, lecun_normal, lecun_uniform, glorot_normal, glorot_uniform, he_normal and he_uniform. They are called as initializer(key, shape, dtype) and are useful when a weight is generated directly instead of through an Equinox layer.

import jax.numpy as jnp
import quantax as qtx

W = qtx.nn.lecun_normal(qtx.get_subkeys(), (16, 8), jnp.float32)

`compute_sign`(kernel, s, output[, neg, bias])	Compute a sign, phase, or cosine structure from a linear form of the spin configuration, `phase = jnp.dot(kernel.flatten(), s.flatten()) + bias`.
`marshall_sign`(s)	Marshall sign rule for bipartite lattices.
`stripe_sign`(s[, alternate_dim])	Stripe sign rule for bipartite lattices.
`neel120_phase`(s)	120 degree Neel phase for triangular lattices.

`ReshapeConv`(dtype)	Reshape the input to the shape suitable for convolutional layers.
`Conv`(in_channels, out_channels, kernel_size, ...)	Convolution layer whose padding is determined by the lattice boundary.
`ConvSymmetrize`([symm])	Symmetrize the output of a convolutional network by translations.
`circular_pad`(x, padding)	Circular padding on all physical dimensions of a feature map.

`input_to_index`(s)	Convert a \(\pm 1\) configuration into a per-unit-cell state index.
`Embedding`(d, sublattice, dtype)	Learned positional encoding of a configuration onto the lattice grid.

`fermion_idx`	Get the indices of occupied fermion sites.
`changed_inds`	Get the indices of the hopping fermions.
`permute_sign`	Get the sign change due to fermion hopping.
`fermion_inverse_sign`	Usual fermion ordering goes from site 0 to site N-1, while QuSpin orders from site N-1 to site 0.
`fermion_reorder_sign`	Compute the sign difference between two fermionic orderings, one having channel dimension first, the other having channel dimension last.