flax

78d85af6 - * Adds `pmap` and `Pmap`. `static_broadcasted_argnums`, `donate_argnums`, and `global_arg_shapes` are not yet supported.

* Adds `pmap` and `Pmap`. `static_broadcasted_argnums`, `donate_argnums`, and `global_arg_shapes` are not yet supported. * Refactors `vmap` (taken from #3949). * Splits the `transforms` module into multiple files. TODO: currently `transforms.transforms` still contains many of the existing transforms that need to be moved to separate files. PiperOrigin-RevId: 641868964