Neural Tangents Versions Save

New Features:

• nt.stax.Sign
• Allow to pass a to_dense function to nt.stax.Aggregate to allow storing the entire graph in a sparse format in GNNs. See #86.
• Support get="ntkgp" in `nt.predict.gp_inference (thanks @bobby-he!). See #93 and https://arxiv.org/abs/2007.05864.

Bugfixes:

• Improve numerical stability of differentiating nonlinearities, and avoid NaNs, notably in nt.stax.Relu. See #88 and #73.
• Allow to pass different test/train kwargs in nt.predict.gradient_descent_mse_ensemble. See #79.

New features:

• Major speedup of the empirical NTK via vmap_axes - please see https://neural-tangents.readthedocs.io/en/latest/neural_tangents.empirical.html and discussion in #30
• Allow to compute maximum theoretical learning rate for a momentum optimizer
• Add an IMDB sentiment analysis example
• Allow pytrees as outputs of functions to linearize/taylorize

Breaking changes:

• Fuse nt.empirical_direct_ntk_fn, nt.empirical_ntk_fn into a single nt.empirical_direct_ntk_fn accepting implementation=1/2 argument (1 - direct, default; 2 - implicit)
• Rename nt.stax.NumericalActivation into nt.stax.ElementwiseNumerical

Minor bugfixes

New layers:

• nt.stax.DotGeneral
• nt.stax.ConvTranspose
• nt.stax.ConvLocal

New features:

• Support arbitrary nested structures of arrays (tuples, lists of lists etc) as inputs and/or outputs
• You can now debug your NT code with config.update('jax_numpy_rank_promotion', 'raise'), as NT and related JAX code no longer rank-promote

Breaking changes:

• Remove nt.stax.GeneralConv (use nt.stax.Conv instead)

• nt.stax.Sigmoid_like
• nt.stax.FanInProd
• Speedup masking
• Bugfixes

• Arbitrary nonlinearity approximation via Gaussian quadrature (stax.NumericalActivation)
• Initial support for Graph Neural Networks (stax.Aggregate)
• Further improvements and new features for attention (stax.GlobalSelfAtttention) from Infinite attention: NNGP and NTK for deep attention networks

• New nonlinearities: a Erf(b x + c), Gelu, RBF
• Numpy 1.19 compatibility
• Bugfixes and documentation improvements

• Huge overhaul of nt.predict and nt.utils.empirical for improved performance and cleaner API - breaking change!

• Support for a sin(b x + c) activation.

• Much improved and expanded documentation and type annotations