It has stax, which is a minimal example of how to build a high level library: https://github.com/google/jax/blob/master/jax/experimental/s...

It seems that the JAX developers are focusing their time on making the core framework better and are leaving the task of building high-level APIs to the community for now. I suspect we'll see a few high-level APIs emerge over the next few months that explore different approaches before the community settles on a particular one.