flax

161b4771 - Tests imagenet example with _ResNet1.

Tests imagenet example with _ResNet1. This reduces train_test.train_and_evaluate time from 122s (61s compilation) to 18s (2s compilation).