First pass at making error_if work with sharding.#342

Draft

patrick-kidger wants to merge 1 commit into

mainfrom

Commits on May 8, 2023

First pass at making error_if work with sharding. Waiting on JAX bugs to be fixed before taking this further. TODO: make sure that things don't get replicated.
patrick-kidger
committed