jax.example_libraries module
jax.example_libraries module#
JAX provides some small, experimental libraries for machine learning. Theselibraries are in part about providing tools and in part about serving asexamples for how to build such libraries using JAX. Each one is only <300 sourcelines of code, so take a look inside and adapt them as you need!
Note
Each mini-library is meant to be aninspiration, but not a prescription.
To serve that purpose, it is best to keep their code samples minimal; so wegenerallywill not merge PRs adding new features. Instead, please send yourlovely pull requests and design ideas to more fully-featured libraries likeHaiku orFlax.