neural_tangents.stax.Hermite

neural_tangents.stax.Hermite(degree)

Hermite polynomials.

Inputs to this layer are assumed to have unit norm, i.e. np.std(x, axis=channel_axis) == 1. The Hermite polynomials are normailized so that the L2 norm w.r.t. standard Gaussian is 1.

Parameters

degree (int) – an integer between 1 and 6.

Return type

Tuple[InitFn, ApplyFn, LayerKernelFn]

Returns

(init_fn, apply_fn, kernel_fn).