[BUG] XLA is incompatible with jax 0.4.29

[pseudo-rnd-thoughts]

Describe the bug

jax.interpreters.xla.backend_specific_translations is deprecated in jax v0.4.29
https://jax.readthedocs.io/en/latest/changelog.html#jax-0-4-29-june-10-2024

This causes the following error when running in xla mode AttributeError: jax.interpreters.xla.backend_specific_translations is deprecated. Register custom primitives via jax.interpreters.mlir instead.

To Reproduce

import envpool

env = envpool.make("Breakout-v5")
env.xla()

    handle, recv, send, step = env.xla()
  File "/lib/python3.10/site-packages/envpool/python/lax.py", line 30, in xla
    _handle, _recv, _send = make_xla(self)
  File "/lib/python3.10/site-packages/envpool/python/xla_template.py", line 124, in make_xla
    methods.append(_make_xla_function(obj, handle, name, specs, capsules))
  File "/lib/python3.10/site-packages/envpool/python/xla_template.py", line 94, in _make_xla_function
    xla.backend_specific_translations["cpu"][prim] = partial(
  File "/lib/python3.10/site-packages/jax/_src/deprecations.py", line 52, in getattr
    raise AttributeError(message)
AttributeError: jax.interpreters.xla.backend_specific_translations is deprecated. Register custom primitives via jax.interpreters.mlir instead.

Expected behavior

xla function is created

System info

import envpool, numpy, sys
print(envpool.__version__, numpy.__version__, sys.version, sys.platform)
> 0.8.4 1.26.4 3.10.12 (main, Jul 29 2024, 16:56:48) [GCC 11.4.0] linux
print(jax.__version__)
> 0.4.29

Reason and Possible fixes

According to the error message then we should use Register custom primitives via jax.interpreters.mlir instead

Checklist

• I have checked that there is no similar issue in the repo (required)
• I have read the documentation (required)
• I have provided a minimal working example to reproduce the bug (required)

[pseudo-rnd-thoughts]

Closed by accident

[Trinkle23897]

sorry about that, I'll try fixing the ci if I have time

[Trinkle23897]

Thanks for the report. This is fixed on main now via PR #336.

A few concrete details:

• the XLA/JAX path was updated as part of the JAX upgrade work on main
• the repo now targets jax[cpu]>=0.9.2,<0.10
• the Atari XLA coverage is back to running in CI rather than being skipped (test_xla_api and test_xla_correctness in envpool/atari/atari_envpool_test.py)

One caveat: this fix is on main, but it has not been cut into a GitHub/PyPI release yet. The latest release is still v0.8.4, which predates the fix.

Also note that current main now targets Python >= 3.11.

PR: #336