Update optax requirement from <0.0.10,>=0.0.2 to >=0.0.2,<0.1.1 (#991)

* Update optax requirement from <0.0.10,>=0.0.2 to >=0.0.2,<0.1.1

Updates the requirements on [optax](https://github.com/deepmind/optax) to permit the latest version.
- [Release notes](https://github.com/deepmind/optax/releases)
- [Commits](https://github.com/deepmind/optax/compare/v0.0.5...v0.1.0

)

---
updated-dependencies:
- dependency-name: optax
  dependency-type: direct:production
...

Signed-off-by: dependabot[bot] <support@github.com>

* Update setup.py

Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>
Co-authored-by: Filippo Vicentini <filippovicentini@gmail.com>