v0.6.0

New features:

• nt.empirical:

  • New implementation=3 for nt.empirical, allowing to often speed-up or reduce the memory of the empirical NTK by orders of magnitude. Please see our ICML2022 paper Fast Finite Width Neural Tangent Kernel, new empirical NTK examples, and visit us on Thursday at ICML in-person!
  • New experimental prototype of using our empirical NTK implementations in Tensorflow via nt.experimental.empirical_ntk_fn_tf.
  • Make nt.empircial work with arbitrary pytrees.

• nt.stax:

  • nt.stax.Index layer mimicking numpy indexing.
  • Expose the layer scaling parameter s in parameterization="standard".

Improvements:

• Slightly lower memory usage in batching.
• Many improvements to documentation and type annotations.
• Simplify test specifications and avoid relying on JAX testing utilities.

Bugfixes:

• Make nt.batch use the correct (local) number of devices in a multi-host setting (thanks @jglaser).

Breaking changes:

• If you've been using nt.empirical or nt.monte_carlo with mask_constant (empirical kernels with masking, for variable-lengths sequences), you now need to pass nt.empirical_kernel_fn(stax.unmask_fn(apply_fn)) instead of nt.empirical_kernel_fn(apply_fn). This is done in order to decouple nt.empirical and nt.stax, and the under-the-hood unmasking happening in nt.empirical is now delegated to stax.unmask_fn.