How do I save an optimizer state of JAX trained model?

Question

I am playing with the mnist_vae example and can\'t figure out how to properly save/load weights of the trained model.

enc_init_rng, dec_init_rng = random.spli