[Feature Request] JAX training

[njzjz]

Summary

Implement JAX training.

Detailed Description

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

Further Information, Files, and Links

No response