flax
permit passing mesh or rules to spmd primitives
#2939

Merged

permit passing mesh or rules to spmd primitives #2939

copybara-service merged 4 commits into main from improve_with_logical_constraint

cgarciae requested a review from

levskaya 2 years ago

cgarciae requested a review from

IvyZX 2 years ago

IvyZX commented on 2023-03-09

permit passing mesh or rules to spmd primitives

97668a91

cgarciae force pushed to 97668a91 2 years ago

IvyZX approved these changes on 2023-03-13

add rules and mesh docstrings

d053913e

ignore FutureWarning from jax

d146765a

fix mypy errors

5b6579d2

IvyZX added pull ready

copybara-service merged 351248c8 into main 2 years ago

copybara-service deleted the improve_with_logical_constraint branch 2 years ago

Reviewers

IvyZX

levskaya

Assignees

No one assigned

Labels

pull ready

Milestone

No milestone