jax

9fd69a04 - Replace uses of ExecutePerReplica with ExecuteOnLocalDevices. (#2394)

Replace uses of ExecutePerReplica with ExecuteOnLocalDevices. (#2394) ExecutePerReplica is deprecated, and ExecuteOnLocalDevices is now available via the minimum jaxlib version.