flax

018725f3 - Merge pull request #2812 from nouiz:imagenet_benchmark