[gdahl]

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.