Optax

What's Changed
* Tweak the meta-learning example from the docs by copybara-service in https://github.com/deepmind/optax/pull/233
* Fix small bugs in metalearning example. by copybara-service in https://github.com/deepmind/optax/pull/236
* Do not reuse mini-batches between epochs in DPSGD example. by copybara-service in https://github.com/deepmind/optax/pull/230
* Make the version of typing_extensions less constrained. by copybara-service in https://github.com/deepmind/optax/pull/238
* [JAX] move example libraries from `jax.experimental` to `jax.example_libraries` by copybara-service in https://github.com/deepmind/optax/pull/200
* Export ScaleByBeliefState by NeilGirdhar in https://github.com/deepmind/optax/pull/239
* MultiStep optimizer: align parameter naming and type annotations of update function with signature of GradientTransform.update. by copybara-service in https://github.com/deepmind/optax/pull/243
* Fix imports of datasets in examples folder. by copybara-service in https://github.com/deepmind/optax/pull/242
* Enable example tests on github. Fix the bugs that were uncovered. by copybara-service in https://github.com/deepmind/optax/pull/244
* Formatting. by copybara-service in https://github.com/deepmind/optax/pull/249
* Add test for multi steps wrapper, verifying that the aggregated gradient is the mean of the input gradients. by copybara-service in https://github.com/deepmind/optax/pull/255
* MultiStep optimizer wrapper: replace naive streaming average gradient implementation with numerically stabler one. by copybara-service in https://github.com/deepmind/optax/pull/254
* Added ord, axis, and keepdims args to safe_norm by copybara-service in https://github.com/deepmind/optax/pull/252
* Add badges and include RTD build into CI tests. by copybara-service in https://github.com/deepmind/optax/pull/256
* Write a clearer doc-string for GradientTransformation by copybara-service in https://github.com/deepmind/optax/pull/257
* Refactor clipping.py by copybara-service in https://github.com/deepmind/optax/pull/260
* Implement split real norm by wdphy16 in https://github.com/deepmind/optax/pull/241
* Monkey-patch sphinx to output correct type annotations for the most common cases (e.g. params, opt state) in the documentation. by copybara-service in https://github.com/deepmind/optax/pull/266
* Improve docs by copybara-service in https://github.com/deepmind/optax/pull/268
* Implement stateless wrapper. by n2cholas in https://github.com/deepmind/optax/pull/246
* Replace _ with params to ensure you can always call init with named args. by copybara-service in https://github.com/deepmind/optax/pull/270
* Improve docs. by copybara-service in https://github.com/deepmind/optax/pull/269
* Add missing ` in two places. by copybara-service in https://github.com/deepmind/optax/pull/273
* Add option to cache examples datasets after pre-processing. by copybara-service in https://github.com/deepmind/optax/pull/272
* Fix an error in README.md rendering. by copybara-service in https://github.com/deepmind/optax/pull/275
* Remove the old venv directory before testing the package. by copybara-service in https://github.com/deepmind/optax/pull/289
* Fix Yogi optimizer by wdphy16 in https://github.com/deepmind/optax/pull/288
* Bump ipython from 7.16.1 to 7.16.3 in /requirements by dependabot in https://github.com/deepmind/optax/pull/286
* Clarifies `optax.adamw(mask)` parameter. by copybara-service in https://github.com/deepmind/optax/pull/284
* Fix the link to the complex-valued optim proposal in RTD. by copybara-service in https://github.com/deepmind/optax/pull/295
* Implement complex norm in optimizers by wdphy16 in https://github.com/deepmind/optax/pull/279
* Change add_noise to match the target variance by scaling by its sqrt. by Rupt in https://github.com/deepmind/optax/pull/294
* Minor tweaks to the optax documentation. by copybara-service in https://github.com/deepmind/optax/pull/297
* Bump version to 0.1.1 from 0.1.0 by copybara-service in https://github.com/deepmind/optax/pull/298

New Contributors
* wdphy16 made their first contribution in https://github.com/deepmind/optax/pull/241
* dependabot made their first contribution in https://github.com/deepmind/optax/pull/286
* Rupt made their first contribution in https://github.com/deepmind/optax/pull/294

**Full Changelog**: https://github.com/deepmind/optax/compare/v0.1.0...v0.1.1

Support for Python 3.6 has been dropped following the [JAX deprecation policy](https://jax.readthedocs.io/en/latest/deprecation.html). Please upgrade to a supported Python version.

**Full Changelog**: https://github.com/deepmind/optax/compare/v0.0.91...v0.1.0

It is the latest version compatible with Python 3.6. See 222 for more details.

Added:
* multi_transform
* LARS optimiser
* AdaFactor optimiser

Fixed:
* masked wrapper for empty params nest
* label type in loss.py
* checks in loss.py
* MultiSteps arguments handling
* Chex asserts in constrain_test.py

* Added `clipping_by_block_rms`
* Added `sgdr_schedule`
* Fixed inconsistency in ema's outputs
* Added `linear_schedule`

Changelog

[Full Changelog](https://github.com/deepmind/optax/compare/v0.0.5...HEAD)

**Closed issues:**

- Documentation [\3](https://github.com/deepmind/optax/issues/3)

**Merged pull requests:**

- Fix section headers in docs [\79](https://github.com/deepmind/optax/pull/79) ([KristianHolsheimer](https://github.com/KristianHolsheimer))
- Define MANIFEST.in to include requirements.txt in pypi bundle. Increment the version. [\76](https://github.com/deepmind/optax/pull/76) ([copybara-service[bot]](https://github.com/apps/copybara-service))
- Fix RTD build error. [\74](https://github.com/deepmind/optax/pull/74) ([copybara-service[bot]](https://github.com/apps/copybara-service))
- Improve doc-strings and formatting of popular optimiser instances - Part 1 [\71](https://github.com/deepmind/optax/pull/71) ([copybara-service[bot]](https://github.com/apps/copybara-service))
- Set up sphinix-built documentation for RTD to build and serve documentation from the public Github. [\61](https://github.com/deepmind/optax/pull/61) ([copybara-service[bot]](https://github.com/apps/copybara-service))



\* *This Changelog was automatically generated by [github_changelog_generator](https://github.com/github-changelog-generator/github-changelog-generator)*