v0.6.1

New features:

• nt.stax:

  • New nonlinearities (Gabor, Polynomial, Monomial, RectifiedMonomial) following the release of Fast Neural Kernel Embeddings for General Activations.

• nt.empirical:

  • An efficient NTK-vector product function nt.empirical_ntk_vp_fn (without instantiating the NTK).

Improvements:

• Make nt.stax.Hermite support arbitrary degrees.
• Various documentation improvements.