How to resume language modeling training with flax?

I know that Trainer has an easy way to resume training. Nonetheless, the mlm script for flax does not contain this functionality. How the training could be resumed with flax in mlm? (restoring weights, optimizer state, etc)