flax

c2a63461 - Merge pull request #298 from danielsuo/dsuo/update-tensorboard-tests