Testing Jax model training. After a few iterations (still in the first epoch), testing model save

Testing JAX training with TPU VM, After 1 Epoch

initial commit