Title: Support jax 0.4.27 in CI tests
Author: albertvillanova (MEMBER)
State: closed
Created: 2024-05-08T09:19:37Z

Support jax 0.4.27 in CI tests by using jax Array `devices` method instead of `device` (which no longer exists).

Fix #6884.

Comments:
The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/datasets/pr_6885). All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

<details>
<summary>Show benchmarks</summary>

PyArrow==8.0.0

<details>
<summary>Show updated benchmarks!</summary>

### Benchmark: benchmark_array_xd.json

| metric | read_batch_formatted_as_numpy after write_array2d | read_batch_formatted_as_numpy after write_flattened_sequence | read_batch_formatted_as_numpy after write_nested_sequence | read_batch_unformated after write_array2d | read_batch_unformated after write_flattened_sequence | read_batch_unformated after write_nested_sequence | read_col_formatted_as_numpy after write_array2d | read_col_formatted_as_numpy after write_flattened_sequence | read_col_formatted_as_numpy after write_nested_sequence | read_col_unformated after write_array2d | read_col_unformated after write_flattened_sequence | read_col_unformated after write_nested_sequence | read_formatted_as_numpy after write_array2d | read_formatted_as_numpy after write_flattened_sequence | read_formatted_as_numpy after write_nested_sequence | read_unformated after write_array2d | read_unformated after write_flattened_sequence | read_unformated after write_nested_sequence | write_array2d | write_flattened_sequence | write_nested_sequence |
|--------|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|
| new / old (diff) | 0.005232 / 0.011353 (-0.006121) | 0.003749 / 0.011008 (-0.007260) | 0.063451 / 0.038508 (0.024943) | 0.031164 / 0.023109 (0.008055) | 0.252024 / 0.275898 (-0.023874) | 0.274479 / 0.323480 (-0.049001) | 0.003238 / 0.007986 (-0.004748) | 0.002668 / 0.004328 (-0.001660) | 0.049570 / 0.004250 (0.045320) | 0.046159 / 0.037052 (0.009107) | 0.273416 / 0.258489 (0.014927) | 0.299064 / 0.293841 (0.005223) | 0.027758 / 0.128546 (-0.100788) | 0.010702 / 0.075646 (-0.064944) | 0.207244 / 0.419271 (-0.212028) | 0.036139 / 0.043533 (-0.007394) | 0.249966 / 0.255139 (-0.005173) | 0.270685 / 0.283200 (-0.012515) | 0.019938 / 0.141683 (-0.121745) | 1.133642 / 1.452155 (-0.318512) | 1.170712 / 1.492716 (-0.322004) |

### Benchmark: benchmark_getitem\_100B.json

| metric | get_batch_of\_1024\_random_rows | get_batch_of\_1024\_rows | get_first_row | get_last_row |
|--------|---|---|---|---|
| new / old (diff) | 0.098352 / 0.018006 (0.080346) | 0.310738 / 0.000490 (0.310248) | 0.000225 / 0.000200 (0.000025) | 0.000044 / 0.000054 (-0.000010) |

### Benchmark: benchmark_indices_mapping.json

| metric | select | shard | shuffle | sort | train_test_split |
|--------|---|---|---|---|---|
| new / old (diff) | 0.018151 / 0.037411 (-0.019261) | 0.061169 / 0.014526 (0.046644) | 0.073275 / 0.176557 (-0.103281) | 0.120320 / 0.737135 (-0.616815) | 0.083945 / 0.296338 (-0.212394) |

### Benchmark: benchmark_iterating.json

| metric | read 5000 | read 50000 | read_batch 50000 10 | read_batch 50000 100 | read_batch 50000 1000 | read_formatted numpy 5000 | read_formatted pandas 5000 | read_formatted tensorflow 5000 | read_formatted torch 5000 | read_formatted_batch numpy 5000 10 | read_formatted_batch numpy 5000 1000 | shuffled read 5000 | shuffled read 50000 | shuffled read_batch 50000 10 | shuffled read_batch 50000 100 | shuffled read_batch 50000 1000 | shuffled read_formatted numpy 5000 | shuffled read_formatted_batch numpy 5000 10 | shuffled read_formatted_batch numpy 5000 1000 |
|--------|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|
| new / old (diff) | 0.283285 / 0.215209 (0.068075) | 2.766129 / 2.077655 (0.688475) | 1.477831 / 1.504120 (-0.026289) | 1.363365 / 1.541195 (-0.177830) | 1.402081 / 1.468490 (-0.066409) | 0.554100 / 4.584777 (-4.030677) | 2.374885 / 3.745712 (-1.370827) | 2.866260 / 5.269862 (-2.403601) | 1.775109 / 4.565676 (-2.790567) | 0.062416 / 0.424275 (-0.361859) | 0.005490 / 0.007607 (-0.002117) | 0.379293 / 0.226044 (0.153248) | 3.330534 / 2.268929 (1.061606) | 1.881648 / 55.444624 (-53.562977) | 1.549847 / 6.876477 (-5.326629) | 1.660350 / 2.142072 (-0.481722) | 0.631013 / 4.805227 (-4.174214) | 0.116646 / 6.500664 (-6.384018) | 0.042977 / 0.075469 (-0.032492) |

### Benchmark: benchmark_map_filter.json

| metric | filter | map fast-tokenizer batched | map identity | map identity batched | map no-op batched | map no-op batched numpy | map no-op batched pandas | map no-op batched pytorch | map no-op batched tensorflow |
|--------|---|---|---|---|---|---|---|---|---|
| new / old (diff) | 0.996102 / 1.841788 (-0.845685) | 12.079143 / 8.074308 (4.004835) | 9.903568 / 10.191392 (-0.287824) | 0.141447 / 0.680424 (-0.538976) | 0.014115 / 0.534201 (-0.520086) | 0.287576 / 0.579283 (-0.291707) | 0.262951 / 0.434364 (-0.171413) | 0.325167 / 0.540337 (-0.215170) | 0.425780 / 1.386936 (-0.961156) |

</details>
PyArrow==latest

<details>
<summary>Show updated benchmarks!</summary>

### Benchmark: benchmark_array_xd.json

| metric | read_batch_formatted_as_numpy after write_array2d | read_batch_formatted_as_numpy after write_flattened_sequence | read_batch_formatted_as_numpy after write_nested_sequence | read_batch_unformated after write_array2d | read_batch_unformated after write_flattened_sequence | read_batch_unformated after write_nested_sequence | read_col_formatted_as_numpy after write_array2d | read_col_formatted_as_numpy after write_flattened_sequence | read_col_formatted_as_numpy after write_nested_sequence | read_col_unformated after write_array2d | read_col_unformated after write_flattened_sequence | read_col_unformated after write_nested_sequence | read_formatted_as_numpy after write_array2d | read_formatted_as_numpy after write_flattened_sequence | read_formatted_as_numpy after write_nested_sequence | read_unformated after write_array2d | read_unformated after write_flattened_sequence | read_unformated after write_nested_sequence | write_array2d | write_flattened_sequence | write_nested_sequence |
|--------|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|
| new / old (diff) | 0.005213 / 0.011353 (-0.006139) | 0.003686 / 0.011008 (-0.007322) | 0.049963 / 0.038508 (0.011455) | 0.030635 / 0.023109 (0.007525) | 0.263992 / 0.275898 (-0.011906) | 0.289960 / 0.323480 (-0.033520) | 0.004281 / 0.007986 (-0.003704) | 0.002709 / 0.004328 (-0.001619) | 0.049147 / 0.004250 (0.044897) | 0.041036 / 0.037052 (0.003984) | 0.277621 / 0.258489 (0.019132) | 0.305689 / 0.293841 (0.011848) | 0.029342 / 0.128546 (-0.099205) | 0.010350 / 0.075646 (-0.065296) | 0.058221 / 0.419271 (-0.361051) | 0.033774 / 0.043533 (-0.009759) | 0.266163 / 0.255139 (0.011024) | 0.286866 / 0.283200 (0.003666) | 0.018463 / 0.141683 (-0.123219) | 1.136930 / 1.452155 (-0.315225) | 1.193974 / 1.492716 (-0.298742) |

### Benchmark: benchmark_getitem\_100B.json

| metric | get_batch_of\_1024\_random_rows | get_batch_of\_1024\_rows | get_first_row | get_last_row |
|--------|---|---|---|---|
| new / old (diff) | 0.106787 / 0.018006 (0.088781) | 0.304229 / 0.000490 (0.303740) | 0.000209 / 0.000200 (0.000009) | 0.000043 / 0.000054 (-0.000012) |

### Benchmark: benchmark_indices_mapping.json

| metric | select | shard | shuffle | sort | train_test_split |
|--------|---|---|---|---|---|
| new / old (diff) | 0.022066 / 0.037411 (-0.015346) | 0.075510 / 0.014526 (0.060984) | 0.087273 / 0.176557 (-0.089284) | 0.128050 / 0.737135 (-0.609085) | 0.090492 / 0.296338 (-0.205847) |

### Benchmark: benchmark_iterating.json

| metric | read 5000 | read 50000 | read_batch 50000 10 | read_batch 50000 100 | read_batch 50000 1000 | read_formatted numpy 5000 | read_formatted pandas 5000 | read_formatted tensorflow 5000 | read_formatted torch 5000 | read_formatted_batch numpy 5000 10 | read_formatted_batch numpy 5000 1000 | shuffled read 5000 | shuffled read 50000 | shuffled read_batch 50000 10 | shuffled read_batch 50000 100 | shuffled read_batch 50000 1000 | shuffled read_formatted numpy 5000 | shuffled read_formatted_batch numpy 5000 10 | shuffled read_formatted_batch numpy 5000 1000 |
|--------|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|
| new / old (diff) | 0.299034 / 0.215209 (0.083825) | 2.899115 / 2.077655 (0.821461) | 1.625169 / 1.504120 (0.121049) | 1.456491 / 1.541195 (-0.084703) | 1.433063 / 1.468490 (-0.035427) | 0.565416 / 4.584777 (-4.019361) | 0.979298 / 3.745712 (-2.766415) | 2.748965 / 5.269862 (-2.520897) | 1.738671 / 4.565676 (-2.827005) | 0.062869 / 0.424275 (-0.361407) | 0.005001 / 0.007607 (-0.002606) | 0.348534 / 0.226044 (0.122489) | 3.437791 / 2.268929 (1.168862) | 1.896804 / 55.444624 (-53.547821) | 1.658544 / 6.876477 (-5.217933) | 1.649106 / 2.142072 (-0.492966) | 0.653791 / 4.805227 (-4.151436) | 0.125522 / 6.500664 (-6.375142) | 0.051260 / 0.075469 (-0.024209) |

### Benchmark: benchmark_map_filter.json

| metric | filter | map fast-tokenizer batched | map identity | map identity batched | map no-op batched | map no-op batched numpy | map no-op batched pandas | map no-op batched pytorch | map no-op batched tensorflow |
|--------|---|---|---|---|---|---|---|---|---|
| new / old (diff) | 1.025170 / 1.841788 (-0.816617) | 12.247968 / 8.074308 (4.173660) | 9.863777 / 10.191392 (-0.327615) | 0.140498 / 0.680424 (-0.539926) | 0.015158 / 0.534201 (-0.519043) | 0.288210 / 0.579283 (-0.291073) | 0.128207 / 0.434364 (-0.306157) | 0.398735 / 0.540337 (-0.141603) | 0.418217 / 1.386936 (-0.968719) |

</details>
</details>

![](https://cml.dev/watermark.png#871eabc7b23c27d677bc06ae2cc1ec3a2a04b10f "CML watermark")