flax
* Adds `pmap` and `Pmap`. `static_broadcasted_argnums`, `donate_argnums`, and `global_arg_shapes` are not yet supported.
#3978

Merged

* Adds `pmap` and `Pmap`. `static_broadcasted_argnums`, `donate_argnums`, and `global_arg_shapes` are not yet supported. #3978

ghost merged 1 commit into main from test_641196775

copybara-service force pushed 1 year ago

copybara-service force pushed to ec1a5a4d 1 year ago

copybara-service force pushed from ec1a5a4d 1 year ago

copybara-service force pushed 1 year ago

* Adds `pmap` and `Pmap`. `static_broadcasted_argnums`, `donate_argnu…

78d85af6

copybara-service force pushed to 78d85af6 1 year ago

copybara-service closed this 1 year ago

copybara-service deleted the test_641196775 branch 1 year ago

Reviewers

No reviews

Assignees

No one assigned

Labels

None yet

Milestone

No milestone