Documentation request: saving a model and loading after training

Could you please add to the documentation what is the way to save models when using MLX - after training is complete final model for inferencing.

Perhaps even add sample code to one of the mlx_examples e.g transformer_lm. How to save for checkpoints would be useful too.

I see multiple methods to save (e.g mx.savez(), model.save_weights()) and unclear whats the best way that saves all the required state and the corresponding methods to load it back from disk.

Start with the transformer_lm example and compare the documented save methods, including mx.savez() and model.save_weights(). Update the documentation with a clear recommendation and matching load procedure for a final inference model and checkpoints, including sample code where appropriate.