Summary

Implement JAX training.

Detailed Description

Need to implement loss, data, and data stat first.

Further Information, Files, and Links

No response