jax

4bc95df7 - [pmap] Remove the config state `jax_pmap_shmap_merge`.

[pmap] Remove the config state `jax_pmap_shmap_merge`. - Delete `disable_pmap_shmap_merge`. - Delete tests for `jax.config.jax_pmap_shmap_merge=False`. - Minor formatting to pass linting. PiperOrigin-RevId: 886201411