text: Title: feature request: please wrap optax optimizers
  Body: The [optax library from Deepmind](https://github.com/deepmind/optax) is the future for jax ecosystem (e.g. it will be adopted by Flax team), and will likely replace [jax.experimental.optimizers](https://jax.readthedocs.io/en/latest/jax.experimental.optimizers.html). It would be good to update https://num.pyro.ai/en/stable/optimizers.html accordingly.
  