Migrate to optax

Created by: epignatelli

Migrate from using jax.experimental.optimisers to deepmind's optax