Constrain optax to 0.1.x

[IamJeffG]

This assists downstream users of scenic (example: see https://github.com/google-research/perch/issues/613) by constraining scenic's dependency on optax to

1. a specific minor version: as optax is still a 0.y.z project we might assume that the 0.1.z to 0.2.0 change contains breaking changes.
2. actual releases of optax, not a github URL

Disclaimer: I don't know this history behind the extant dependency specification, e.bg. maybe there was a reason scenic needs to rely on bleeding-edge pre-release code. Would the scenic authors be amenable to my change?

[google-cla[bot]]

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.