`ArrayTest.testThreadsafeIndexing` fails under `pmap_shmap_merge=True`.

[danielsuo]

Description

I think this is how to run in open source?

JAX_PMAP_SHMAP_MERGE=1 pytest tests/pmap_test.py::ArrayTest::testThreadsafeIndexing

System info (python version, jaxlib version, accelerator, etc.)

jax:    0.7.2
jaxlib: 0.7.2
numpy:  2.3.2
python: 3.12.11 (stable, redacted, redacted) [Clang 9999.0.0 (fc44a4fcd3c54be927c15ddd9211aca1501633e7)]
device info: cpu-8, 8 local devices"
process_count: 1
platform: uname_result(system='Linux', node='4555d52793ffea74-1577720635f.borgtask.google.com', release='6.6.65-smp-1300.166.0.0', version='#1 [v6.6.65-1300.166.0.0] SMP @1755135593', machine='x86_64')