flax

2204c82d - [jax:benchmarks] Add tracing/lowering benchmarks for a few flax examples.