CI is broken after jax-0.4.27 release: AttributeError: 'jaxlib.xla_extension.DeviceList' object has no attribute 'device'

[albertvillanova]

After jax-0.4.27 release (https://github.com/google/jax/releases/tag/jax-v0.4.27), our CI is broken with the error:

AttributeError: 'jaxlib.xla_extension.DeviceList' object has no attribute 'device'. Did you mean: 'devices'?

See: https://github.com/huggingface/datasets/actions/runs/8997488610/job/24715736153

___________________ FormatterTest.test_jax_formatter_device ____________________
[gw1] linux -- Python 3.10.14 /opt/hostedtoolcache/Python/3.10.14/x64/bin/python

self = <tests.test_formatting.FormatterTest testMethod=test_jax_formatter_device>

    @require_jax
    def test_jax_formatter_device(self):
        import jax
    
        from datasets.formatting import JaxFormatter
    
        pa_table = self._create_dummy_table()
        device = jax.devices()[0]
        formatter = JaxFormatter(device=str(device))
        row = formatter.format_row(pa_table)
>       assert row["a"].device() == device
E       AttributeError: 'jaxlib.xla_extension.DeviceList' object has no attribute 'device'. Did you mean: 'devices'?

tests/test_formatting.py:630: AttributeError