flax

0c4e35a8 - Fix nits in testing/benchmark.py