flax
Fix `get_partition_spec` on replicated array
#3013

Merged

Fix `get_partition_spec` on replicated array #3013

copybara-service merged 1 commit into google:main from IvyZX:pjit

IvyZX requested a review from

cgarciae 2 years ago

cgarciae commented on 2023-04-05

IvyZX force pushed 2 years ago

IvyZX force pushed 2 years ago

fix pspec on replicated array + add sharding utils

efe3d57c

IvyZX force pushed to efe3d57c 2 years ago

cgarciae approved these changes on 2023-04-06

cgarciae added pull ready

copybara-service merged 332b7d04 into main 2 years ago

IvyZX deleted the pjit branch 321 days ago

Reviewers

cgarciae

Assignees

No one assigned

Labels

pull ready

Milestone

No milestone