Check out the Flax ResNet50 example: https://github.com/google/flax/tree/master/...

		gdahl on April 8, 2020 \| parent \| context \| favorite \| on: Jax – Composable transformations of Python and Num... 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.