flax

Add compute_flops and compute_vjp_flops options to `nnx.tabulate`

#4948

Merged