nn

Modules

Sequential(layers[, holomorphic])	A sequence of equinox.Module applied in order similar to Sequential in Equinox.
NoGradLayer()	The layer in which the pytree leaves are not considered as differentiable parameters in Quantax computations.
filter_grad(fun, *[, has_aux])	Creates a function that computes the gradient of fun similar to equinox.filter_grad.
filter_vjp(fun, *primals[, has_aux])	Like equinox.filter_vjp.

Activation function

Scale(scale)	Rescale the input \(f(x) = x * \mathrm{scale}\)
ScaleFn(fn, features, scaling, dtype)	Apply a function to a rescaled input \(f(x) = fn(x * \mathrm{scale})\).
Theta0Layer()	The activation layer with output \(f(x) = g(x) * \exp(\theta_0)\).
SinhShift()	\(f(x) = (\sinh(x) + 1) \exp(\theta_0)\)
Prod()	\(f(x) = \exp(\theta_0) \prod x\)
Exp()	\(f(x) = \exp(x + \theta_0)\)
pair_cpl(x)	Make a real input complex by splitting it into two parts, one taken as the real part and the other the imaginary part.

Initializers

apply_lecun_normal(key, net)	Apply the Lecun normal initializer to the weights of the layer.
apply_he_normal(key, net)	Apply the He normal initializer to the weights of the layer.

NQS layers

ReshapeConv(dtype)	Reshape the input to the shape suitable for convolutional layers.
ConvSymmetrize([symm])	Symmetrize the output of a convolutional network according to the given symmetry.