flax

b26058c0 - Revert "Fix bug in graph overhead benchmark"