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.
cgarciae
commented
3 weeks ago
What does this PR do?
pmapandPmap.static_broadcasted_argnums,donate_argnums, andglobal_arg_shapesare not yet supported.vmap(taken from #3949).transformsmodule into multiple files. TODO: currentlytransforms.transformsstill contains many of the existing transforms that need to be moved to separate files.