bax (pypi): A flexible trainer interface for Jax and Haiku.
Keywords: boilerplate|dm-haiku|jax|trainer