Fix code rendering in optimizers documentation (#4296)
Fix code rendering in optimizers documentation (#4296)

* Fix code rendering in optimizers documentation

* Fix misnamed variable
google/jax,google/jax,google/jax,google/jax,tensorflow/probability,tensorflow/probability