[Feature Request] JAX training #4428
Description
Summary
Implement JAX training.
Detailed Description
Need to implement loss, data, and data stat first.
Further Information, Files, and Links
No response
{{ message }}
Implement JAX training.
Need to implement loss, data, and data stat first.
No response