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

copybara-service[bot] merged 1 commit intomastergoogle-deepmind/clrs:masterfrom

test_888582443google-deepmind/clrs:test_888582443Copy head branch name to clipboard