Neural-tangents

Maintenance release:
* [refactoring and minor improvements](https://github.com/google/neural-tangents/commit/3c3dc9fb0cdce9d8e6ada76ae5d11b6cec060875)
* Support and require
* [JAX >= 0.4.16](https://github.com/google/neural-tangents/commit/ad47437a31a578912cef764496d889f48362b9de)
* [Tensorflow >= 2.15](https://github.com/google/neural-tangents/commit/429dc13617bd1b7e71f5f362bd4b394c1176e633)

Improvements:

* [Support Python 3.11](https://github.com/google/neural-tangents/commit/2c38d69372c9cbfdd750df3e70af358da346bd4d)
* [Use modern generic type annotations](https://github.com/google/neural-tangents/commit/2c38d69372c9cbfdd750df3e70af358da346bd4d)
* Various [bugfixes](https://github.com/google/neural-tangents/commit/ed115c751f544c1ad3fa64e550266be13a013916), [compatibility](https://github.com/google/neural-tangents/commit/c5f8eb9c60b3970b5cffe840ac24bad265df4931) and [documentation](https://github.com/google/neural-tangents/commit/11c024eb879c3a62fee21e2d235e6d90539c4e61) improvements

Breaking changes:

* [Drop support for Python 3.8 following JAX and tf2jax version policy](https://github.com/google/neural-tangents/commit/2c38d69372c9cbfdd750df3e70af358da346bd4d)

New features:

* [`nt.stax.repeat` layer allowing fast compilation of very deep networks](https://github.com/google/neural-tangents/commit/fe98c9c047847a4a31f03aa81992d7328eda29e2) (see #168 and thanks jglaser!)
* Add a [Colab notebook](https://github.com/google/neural-tangents/commit/5ffa88ae4bf0257a7e1af017c49c94cb6bdf228d) accompanying [Precise Learning Curves and Higher-Order Scaling Limits for Dot Product Kernel Regression](https://arxiv.org/abs/2205.14846)

Improvements:
* [Support M1 macs](https://github.com/google/neural-tangents/commit/0272225ee56d8878f5bde30f84d58e058d06b348) (see #171)
* [Internal changes to support newer versions of JAX](https://github.com/google/neural-tangents/commit/edb7af53eed4f3c6e9d1d921fe3dad28aff2daa2)

Breaking changes:
* [Drop support for Python 3.7 following JAX and tf2jax version policy](https://github.com/google/neural-tangents/commit/0272225ee56d8878f5bde30f84d58e058d06b348)
* [On GPU, require CuDNN>=8.2 to support JAX 0.4.3](https://github.com/google/neural-tangents/commit/46306b6a28617a0d03f8d586da132d6766c9b301)

New features:

* `nt.stax`:
* [New nonlinearities (`Gabor`, `Polynomial`, `Monomial`, `RectifiedMonomial`)](https://github.com/google/neural-tangents/commit/47ffb1bd48206b3cdbc5f5137abd6c6036e6cc4a) following the release of **[Fast Neural Kernel Embeddings for General Activations](https://arxiv.org/abs/2209.04121)**.

* `nt.empirical`:
* [An efficient NTK-vector product function `nt.empirical_ntk_vp_fn`](https://github.com/google/neural-tangents/commit/628ce0e58905ba2b16d924c2119d661a01856c83) (without instantiating the NTK).

Improvements:
* [Make `nt.stax.Hermite` support arbitrary degrees](https://github.com/google/neural-tangents/commit/47ffb1bd48206b3cdbc5f5137abd6c6036e6cc4a).
* [Various documentation improvements](https://github.com/google/neural-tangents/commit/6019cc8e013cdacfa5acd53d57fc99b25e6165e0).

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](https://arxiv.org/abs/2206.08720), new [empirical NTK examples](https://github.com/google/neural-tangents#colab-notebooks), and [visit us on Thursday at ICML in-person](https://icml.cc/virtual/2022/spotlight/17194)!
* [New experimental prototype of using our empirical NTK implementations in Tensorflow via `nt.experimental.empirical_ntk_fn_tf`](https://github.com/google/neural-tangents/commit/e28971ce224857094dce8766a8127f1c17aa3c5f).
* [Make `nt.empircial` work with arbitrary pytrees](https://github.com/google/neural-tangents/commit/f86dafae0477900cf1b73ea7030ec7220949a03d).

* `nt.stax`:
* [`nt.stax.Index` layer mimicking numpy indexing](https://github.com/google/neural-tangents/commit/6f10d16979d9a69b1ddac31a4f2a3813abb30c78).
* [Expose the layer scaling parameter `s` in `parameterization="standard"`](https://github.com/google/neural-tangents/commit/239cc849cf55d672018bce0e3539e56b1a50870f).

Improvements:
* [Slightly lower memory usage in batching](https://github.com/google/neural-tangents/commit/d065a8c508d9cbac57604815ea5944d874bc8a28).
* [Many improvements to documentation and type annotations](https://github.com/google/neural-tangents/commit/eec18f56562864540f8141372d7df67f9c34373e).
* [Simplify test specifications and avoid relying on JAX testing utilities](https://github.com/google/neural-tangents/commit/a137b9672f855996748ad5a2f79e9ef0ebc0b8a8).

Bugfixes:
* [Make `nt.batch` use the correct (local) number of devices in a multi-host setting](https://github.com/google/neural-tangents/commit/23bea9f10e40a470521c8ca887be0fc6f5f2570f) (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`](https://github.com/google/neural-tangents/commit/f86dafae0477900cf1b73ea7030ec7220949a03d).

Potentially breaking changes:
- [Significant internal refactoring, notably splitting `stax` into multiple sub-modules, and moving implementations into an `_src` folder](https://github.com/google/neural-tangents/commit/4b183cec5e14b5ada427a0c63303f0ebe6b97854). This could break your code if you use internal function like `nt.utils.typing`, `nt.utils.utils`, `nt.utils.Kernel` etc. (public [API](https://neural-tangents.readthedocs.io/en/latest/) 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.](https://github.com/google/neural-tangents/commit/22b3aaeda3f8ff2f516c4725997bf542bae53741)
- [Improve type annotations via typing protocols](https://github.com/google/neural-tangents/commit/b0d489e58228642e65d7cce9abc7c1a0db0526ea)
- [Run pytype in GitHub Actions](https://github.com/google/neural-tangents/commit/b0d489e58228642e65d7cce9abc7c1a0db0526ea)