Maintenance release:

• refactoring and minor improvements
• Support and require
  • JAX >= 0.4.16
  • Tensorflow >= 2.15

Improvements:

• Support Python 3.11
• Use modern generic type annotations
• Various bugfixes, compatibility and documentation improvements

Breaking changes:

• Drop support for Python 3.8 following JAX and tf2jax version policy

New features:

• nt.stax.repeat layer allowing fast compilation of very deep networks (see #168 and thanks @jglaser!)
• Add a Colab notebook accompanying Precise Learning Curves and Higher-Order Scaling Limits for Dot Product Kernel Regression

Improvements:

• Support M1 macs (see #171)
• Internal changes to support newer versions of JAX

Breaking changes:

• Drop support for Python 3.7 following JAX and tf2jax version policy
• On GPU, require CuDNN>=8.2 to support JAX 0.4.3

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.

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.

Potentially breaking changes:

• Significant internal refactoring, notably splitting stax into multiple sub-modules, and moving implementations into an _src folder. This could break your code if you use internal function like nt.utils.typing, nt.utils.utils, nt.utils.Kernel etc. (public API will remain unchanged). This should be easily fixed by updating the imports, e.g. nt.utils -> nt._src.utils.

New features:

• Improve documentation with (sub-)sections and better type annotations.
• Improve type annotations via typing protocols
• Run pytype in GitHub Actions

WARNING:
Our next major release (v0.5.0) will include significant refactoring, and could break your code if you use internal function like nt.utils.typing, nt.utils.utils, nt.utils.Kernel etc. (public API will remain unchanged). This should be easily fixed by updating the imports, e.g. nt.utils -> nt._src.utils.

This release (v0.4.0):

New feature:

• Add continuous integration with GitHub Actions.
• Python 3.10 support.

Improvements:

• Various internal refactoring and tighter tests.

Bugfixes:

• Fix values and gradients of non-differentiable kernel_fn at zero inputs to be consistent with finite-width kernels, and how JAX defines gradients of non-differentiable functions to be the mean sub-gradient, see also #123.
• Fix wrong treatment of b_std=None in the infinite-width limit with parameterization='standard', see also #123.
• Fix a bug in nt.batch when x2 = None and inputs are PyTrees.

Breaking changes:

• Bump requirements to jax==0.3 and frozendict==2.3.

New Features:

• New nonlinearities:
  • stax.Hermite;
  • stax.Exp;
  • stax.Gaussian;
  • stax.ExpNormalized.
• Support and default to b_std=None in stax layers, treated as symbolic zero, i.e. providing same behavior as b_std=0., but without creating a redundant bias array.

Breaking changes:

• Bump requirements to JAX v0.2.25. In consequence, drop CUDA 10 support to prevent #122
• The b_std=None change could be breaking in very rare edge cases. The dummy bias array is replaced with None in this case and might potentially break your serialization routine.

New Features:

• stax.Elementwise - a layer for generic elementwise functions requiring the user to specify only scalar-valued nngp_fn : (cov12, var1, var2) |-> E[fn(x_1) * fn(x_2)]. The NTK computation (thanks to @SiuMath) and vectorization over the underlying Kernel happen automatically under the hood. If you can't derive the nngp_fn for your function, use stax.ElementwiseNumerical. See docs for more details.

Bugfixes:

• Compatibility with JAX 0.2.21.

Full Changelog: v0.3.7...v0.3.8

New Features:

• nt.stax.Cos
• nt.stax.ImageResize
• New implementation implementation="SPARSE" in nt.stax.Aggregate for efficient handling of sparse graphs (see #86, #9)
• Support approximate=True in nt.stax.Gelu

Bugfixes:

• Fix a bug that might alter Kernel requirements
• Fix nt.batch handling of diagonal_axes (see #87)
• Remove the frequent but redundant warning about type conversion in kernel_fn
• Minor fixes to documentation and code clean-up

Breaking changes:

• Parameters initialized by init_fn now follow the setting of JAX_ENABLE_X64 instead of always defaulting to 32-bit (see #112)
• Drop python 3.6 support and add python 3.9 support