Check out the Flax ResNet50 example: https://github.com/google/flax/tree/master/examples/imagenet
It runs about as fast as any of the other popular machine learning frameworks, occasionally faster.
Disclaimer: I work for Google and use JAX, although I'm not on the Jax team.