Neural-tangents

New layers:
* [`nt.stax.DotGeneral`](https://github.com/google/neural-tangents/commit/b582a89600860d331cb91064e3b0075a9e898c89)
* [`nt.stax.ConvTranspose`](https://github.com/google/neural-tangents/commit/780ad0ce22d482bcefd12f4d3390090de7206da5)
* [`nt.stax.ConvLocal`](https://github.com/google/neural-tangents/commit/cb6cf8027ecc72bc4964f7e24f56508f668678d2)

New features:
* [Support arbitrary nested structures of arrays (tuples, lists of lists etc) as inputs and/or outputs](https://github.com/google/neural-tangents/commit/8a627885296715e51d162e326ebcb5796e30fe4b)
* [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](https://github.com/google/neural-tangents/commit/be351c56e0b5ed293527292afdc9df48a21dc064)

Breaking changes:
* [Remove `nt.stax.GeneralConv` (use `nt.stax.Conv` instead)](https://github.com/google/neural-tangents/commit/dcc77da2c5aef24c2dd9baca00edd98dc47408f7)

* [`nt.stax.Sigmoid_like`](https://github.com/google/neural-tangents/commit/de2a61d3097f9691290f2c67d61ed4f345b53bb6)
* [`nt.stax.FanInProd`](https://github.com/google/neural-tangents/commit/484dfa24b95722206134bbc8a2c457988c794822)
* [Speedup masking](https://github.com/google/neural-tangents/commit/74498fc8d106842869dda14cedf391f524d7c8be)
* [Bugfixes](https://github.com/google/neural-tangents/commit/cfcda489548b0e5bf6e499416ac8b68bf4c612dc)

* [Arbitrary nonlinearity approximation via Gaussian quadrature (`stax.NumericalActivation`)](https://github.com/google/neural-tangents/commit/841a33a03081680b6f09ebcdcc7b4e7b12378d22)
* [Initial support for Graph Neural Networks (`stax.Aggregate`)](https://github.com/google/neural-tangents/commit/03be2577e82c35a3ddd95ef478976ead722e039b)
* [Further improvements and new features for attention (`stax.GlobalSelfAtttention`)](https://github.com/google/neural-tangents/commit/253ddeaa4c34d75baac129efd9891e9c27efc5ee) from [Infinite attention: NNGP and NTK for deep attention networks](https://arxiv.org/abs/2006.10540)

* New nonlinearities: [`a Erf(b x + c)`](https://github.com/google/neural-tangents/commit/d7fc2dbb6844096fd382be78a0869de653542db0), [`Gelu`](https://github.com/google/neural-tangents/commit/70b983d5cb437c06baf213595da0a1c57f6a9533), [`RBF`](https://github.com/google/neural-tangents/commit/6be0179a252b4ff5b1fe8ba7426804f582e9ce1c)
* [Numpy `1.19` compatibility](https://github.com/google/neural-tangents/commit/11d70e11c3ddcc11feae65f502ac0c42291b52ce)
* Bugfixes and documentation improvements

* [Huge overhaul of `nt.predict` and `nt.utils.empirical` for improved performance and cleaner API](https://github.com/google/neural-tangents/commit/a76bbb494f19af4f8c9c1a1b0904e91b105f769e) - breaking change!

* [Support for `a sin(b x + c)` activation.](https://github.com/google/neural-tangents/commit/a936d8794c4cefc27359484c01eac77b4a4ddde1)

* [Much improved and expanded documentation and type annotations](https://neural-tangents.readthedocs.io/en/latest/index.html)