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.
pmap
andPmap
.static_broadcasted_argnums
,donate_argnums
, andglobal_arg_shapes
are not yet supported.vmap
(taken from #3949).transforms
module into multiple files. TODO: currentlytransforms.transforms
still contains many of the existing transforms that need to be moved to separate files.