[pmap] In-line definitions of jax.device_put_sharded and jax.device_put_replicated.#440

copybara-service[bot] merged 1 commit intomaingoogle-deepmind/chex:mainfrom

test_888582452google-deepmind/chex:test_888582452Copy head branch name to clipboard