flax

4c7a45a9 - Report forward and backward pass FLOPs of modules and submodules in `linen.Module.tabulate` and `summary.tabulate` (in new `flops` and `vjp_flops` table columns). Pass `compute_flops=True` and/or `compute_vjp_flops=True` to include these columns.

Report forward and backward pass FLOPs of modules and submodules in `linen.Module.tabulate` and `summary.tabulate` (in new `flops` and `vjp_flops` table columns). Pass `compute_flops=True` and/or `compute_vjp_flops=True` to include these columns. Integrated improved argument splitting by @cgarciae in https://github.com/google/flax/pull/3350 Co-authored-by: Cristian Garcia <cgarcia.e88@gmail.com> PiperOrigin-RevId: 569254535