[pmap] Avoid degraded performance under the new jax.pmap.#334

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

test_846723855google/learned_optimization:test_846723855Copy head branch name to clipboard