[pmap] Remove conditional jax.config.jax_pmap_shmap_merge logic as part of last phase of jax.pmap clean-up.#336

copybara-service[bot] wants to merge 1 commit intomaingoogle/learned_optimization:mainfrom

test_864486731google/learned_optimization:test_864486731Copy head branch name to clipboard