neural_tangents.stax.ExpNormalized

neural_tangents.stax.ExpNormalized(gamma=1, shift=- 1, do_clip=False)[source]

Simulates the “Gaussian normalized kernel”.

Source: https://arxiv.org/abs/2003.02237.pdf, page 6.

Parameters

  • gamma (float) – exponent scalar coefficient.

  • shift (float) – shift exponentiated normalized covariance by this much.

  • do_clip (bool) – True to clip normalized covariance, potentially improving accuracy.

Return type

Tuple[InitFn, ApplyFn, LayerKernelFn]

Returns

(init_fn, apply_fn, `kernel_fn).

Raises

NotImplementedError – if finite width apply_fn is called.