[CI] Failed jaxmd test

https://github.com/EnzymeAD/Enzyme-JAX/actions/runs/16713928087/job/47317194253#step:15:87
```
[7,950 / 7,951] 586 / 587 tests, 1 failed; Testing //test:jaxmd; 64s disk-cache, processwrapper-sandbox
INFO: From Testing //test:jaxmd:
==================== Test output for //test:jaxmd:
Running tests under Python 3.11.10: /root/.bazel/sandbox/processwrapper-sandbox/13802/execroot/__main__/bazel-out/k8-opt/bin/test/jaxmd.runfiles/python_x86_64-unknown-linux-gnu/bin/python3
INFO:2025-08-04 10:00:38,599:jax._src.xla_bridge:830: Unable to initialize backend 'tpu': INTERNAL: Failed to open libtpu.so: libtpu.so: cannot open shared object file: No such file or directory
I0804 10:00:38.599670 133472947390272 xla_bridge.py:830] Unable to initialize backend 'tpu': INTERNAL: Failed to open libtpu.so: libtpu.so: cannot open shared object file: No such file or directory
[ RUN ] test_utils.EnzymeJaxTest.test
[ OK ] test_utils.EnzymeJaxTest.test
[ RUN ] JAXMD.test
/root/.bazel/sandbox/processwrapper-sandbox/13802/execroot/__main__/bazel-out/k8-opt/bin/test/jaxmd.runfiles/pypi_jax_md/site-packages/jax_md/space.py:279: DeprecationWarning: invalid escape sequence '\i'
"""Periodic boundary conditions on a parallelepiped.
/root/.bazel/sandbox/processwrapper-sandbox/13802/execroot/__main__/bazel-out/k8-opt/bin/test/jaxmd.runfiles/pypi_dm_haiku/site-packages/haiku/_src/transform.py:341: DeprecationWarning: jax.lib.xla_extension.PjitFunction is deprecated.
COMPILED_FN_TYPES = (jax.lib.xla_extension.PjitFunction,
/root/.bazel/sandbox/processwrapper-sandbox/13802/execroot/__main__/bazel-out/k8-opt/bin/test/jaxmd.runfiles/pypi_dm_haiku/site-packages/haiku/_src/transform.py:342: DeprecationWarning: jax.lib.xla_extension.PmapFunction is deprecated.
jax.lib.xla_extension.PmapFunction) # pytype: disable=name-error
/root/.bazel/sandbox/processwrapper-sandbox/13802/execroot/__main__/bazel-out/k8-opt/bin/test/jaxmd.runfiles/pypi_jax_md/site-packages/jax_md/quantity.py:328: DeprecationWarning: invalid escape sequence '\s'
"""Computes the pair correlation function at a mesh of distances.
/root/.bazel/sandbox/processwrapper-sandbox/13802/execroot/__main__/bazel-out/k8-opt/bin/test/jaxmd.runfiles/pypi_jax_md/site-packages/jax_md/quantity.py:415: DeprecationWarning: invalid escape sequence '\s'
"""Computes the pair correlation function at a mesh of distances.
/root/.bazel/sandbox/processwrapper-sandbox/13802/execroot/__main__/bazel-out/k8-opt/bin/test/jaxmd.runfiles/pypi_jax_md/site-packages/jax_md/quantity.py:625: DeprecationWarning: invalid escape sequence '\D'
"""Computes the phop indicator of rearrangements.
/root/.bazel/sandbox/processwrapper-sandbox/13802/execroot/__main__/bazel-out/k8-opt/bin/test/jaxmd.runfiles/pypi_ml_collections/site-packages/ml_collections/config_dict/config_dict.py:163: DeprecationWarning: invalid escape sequence '\['
index_match = re.match("(.*)\[([0-9]+)\]", key)
/root/.bazel/sandbox/processwrapper-sandbox/13802/execroot/__main__/bazel-out/k8-opt/bin/test/jaxmd.runfiles/pypi_jax_md/site-packages/jax_md/simulate.py:357: DeprecationWarning: invalid escape sequence '\d'
"""Helper function to simulate a Nose-Hoover Chain coupled to a system.
/root/.bazel/sandbox/processwrapper-sandbox/13802/execroot/__main__/bazel-out/k8-opt/bin/test/jaxmd.runfiles/pypi_jax_md/site-packages/jax_md/elasticity.py:15: DeprecationWarning: invalid escape sequence '\l'
"""Code to calculate the elastic modulus tensor for athermal systems.
/root/.bazel/sandbox/processwrapper-sandbox/13802/execroot/__main__/bazel-out/k8-opt/bin/test/jaxmd.runfiles/pypi_jax_md/site-packages/jax_md/rigid_body.py:322: DeprecationWarning: invalid escape sequence '\d'
"""From Miller III et al., S(q) is defined so that \dot q = 1/2S(q)\omega.
/root/.bazel/sandbox/processwrapper-sandbox/13802/execroot/__main__/bazel-out/k8-opt/bin/test/jaxmd.runfiles/pypi_jax_md/site-packages/jax_md/rigid_body.py:333: DeprecationWarning: invalid escape sequence '\p'
"""Convert from the conjugate momentum of a quaternion to angular momentum.
/root/.bazel/sandbox/processwrapper-sandbox/13802/execroot/__main__/bazel-out/k8-opt/bin/test/jaxmd.runfiles/pypi_jax/site-packages/jax/_src/ops/scatter.py:110: FutureWarning: scatter inputs have incompatible types: cannot safely cast value from dtype=int64 to dtype=int32 with jax_numpy_dtype_promotion='standard'. In future JAX releases this will result in an error.
warnings.warn(
/root/.bazel/sandbox/processwrapper-sandbox/13802/execroot/__main__/bazel-out/k8-opt/bin/test/jaxmd.runfiles/pypi_jax/site-packages/jax/_src/pjit.py:368: DeprecationWarning: backend and device argument on jit is deprecated. You can use `jax.device_put(..., jax.local_devices(backend="cpu")[0])` on the inputs to the jitted function to get the same behavior.
warnings.warn(
[ FAILED ] JAXMD.test
======================================================================
FAIL: test (__main__.JAXMD.test)
JAXMD.test
----------------------------------------------------------------------
Traceback (most recent call last):
File "/root/.bazel/sandbox/processwrapper-sandbox/13802/execroot/__main__/bazel-out/k8-opt/bin/test/jaxmd.runfiles/__main__/test/test_utils.py", line 640, in test
self.harness(self.name, self.fn, self.ins, self.dins, self.douts)
File "/root/.bazel/sandbox/processwrapper-sandbox/13802/execroot/__main__/bazel-out/k8-opt/bin/test/jaxmd.runfiles/__main__/test/test_utils.py", line 880, in harness
recursive_check(self, grads, revres, self.tol)
File "/root/.bazel/sandbox/processwrapper-sandbox/13802/execroot/__main__/bazel-out/k8-opt/bin/test/jaxmd.runfiles/__main__/test/test_utils.py", line 597, in recursive_check
recursive_check(tester, g, g_p, tol, pname)
File "/root/.bazel/sandbox/processwrapper-sandbox/13802/execroot/__main__/bazel-out/k8-opt/bin/test/jaxmd.runfiles/__main__/test/test_utils.py", line 592, in recursive_check
tester.assertTrue(legal)
AssertionError: Array(False, dtype=bool) is not true

----------------------------------------------------------------------
Ran 2 tests in 60.431s

FAILED (failures=1)
Created a system of 64000 LJ particles with number density 0.382
jaxmd40 JaXPipe cpu Primal 0.00011168
jaxmd40 JaX cpu Primal 0.00004952
jaxmd40 HLOOpt cpu Primal 0.00009385
jaxmd40 PartOpt cpu Primal 0.00007393
jaxmd40 DefOpt cpu Primal 0.00007072
jaxmd40 IPartOpt cpu Primal 0.00009052
jaxmd40 IDefOpt cpu Primal 0.00009332
jaxmd40 JaX cpu Forward 0.00015202
jaxmd40 JaXPipe cpu PostRev 0.00014457
jaxmd40 JaX cpu BothRev 0.00016812
lhs [[ 126925.71 -113309.5 -94711.98 ]
[-147981.34 -108944.32 74243.33 ]
[ 91526.8 -128449.65 209822.86 ]
...
[ -19123.922 -180616.27 66605.45 ]
[ 130399.055 117100.8 189371.88 ]
[ -50731.336 -67885.46 107413.625]]
rhs [[ 126925.71 -113309.5 -94711.98 ]
[-147981.34 -108944.32 74243.32 ]
[ 91526.8 -128449.65 209822.86 ]
...
[ -19123.922 -180616.27 66605.45 ]
[ 130399.055 117100.8 189371.88 ]
[ -50731.336 -67885.46 107413.625]]
abs [[0. 0. 0. ]
[0. 0. 0.0078125]
[0. 0. 0. ]
...
[0. 0. 0. ]
[0. 0. 0. ]
[0. 0. 0. ]]
eq [[ True True True]
[ True True True]
[ True True True]
...
[ True True True]
[ True True True]
[ True True True]]
max 0.0625
================================================================================
```

Last successful run on 7ea2f5ade138be1d7d3331bd3d6681e62b94c169, first failure observed in 9a9497a0eb6e9506966e77b7f59db65b194d0314 (build in between were skipped), so breakage should have happened somewhere in https://github.com/EnzymeAD/Enzyme-JAX/compare/7ea2f5ade138be1d7d3331bd3d6681e62b94c169...9a9497a0eb6e9506966e77b7f59db65b194d0314. Can't see the current status because of #1288, but opening the ticket to track the issue.