There are also some ready-made models (and projects) developed in JAX. This list will be extremely helpful for someone doing research or some data science task.

Models and projects

JAX

Fourier Feature Networks: Official implementation of

...