flax

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

#4911

Merged