Add JAX device selection when formatting
MEMBER
The code below was throwing a warning:

```python
class JaxFormatter(Formatter[Mapping, "jax.Array", Mapping]):
    def __init__(self, features=None, device=None, **jnp_array_kwargs):
        super().__init__(features=features)
        import jax
        from jaxlib.xla_extension import Device
        
        self.device = (
            device if isinstance(device, Device) else jax.devices()[0]
        )
        self.jnp_array_kwargs = jnp_array_kwargs

    ...

    def _tensorize(self, value):
        ...

        with jax.default_device(self.device):
            # calling jnp.array on a np.ndarray does copy the data
            # see https://github.com/google/jax/issues/4486
            return jnp.array(value, **{**default_dtype, **self.jnp_array_kwargs})
```

When providing `device` via param:

```python
from datasets import Dataset
import jax

ds = Dataset.from_dict({"a": [1, 2, 3], "b": [4, 5, 6]})
ds = ds.with_format("jax", device=jax.devices()[0])
print(ds[0])
```

Producing the following warning:

```
WARNING:datasets.fingerprint:Parameter 'device'=TFRT_CPU_0 of the transform datasets.arrow_dataset.Dataset.set_format couldn't be hashed properly, a random hash was used instead. Make sure your transforms and parameters are serializable with pickle or dill for the dataset fingerprinting and caching to work. If you reuse this transform, the caching mechanism will consider it to be different from the previous calls and recompute everything. This warning is only showed once. Subsequent hashing failures won't be showed.
```

That's why I decided to map all the available devices, and assign their string representation e.g. `TFRT_CPU_0` to `self.device` instead of `jaxlib.xla_extension.Device`, so that the value of the param `device` is washable. So on, the code that remains at the end is:

```python
class JaxFormatter(Formatter[Mapping, "jax.Array", Mapping]):
    def __init__(self, features=None, device=None, **jnp_array_kwargs):
        super().__init__(features=features)
        import jax
        from jaxlib.xla_client import Device

        self.device_mapping = self._map_devices_to_str()
        self.device = (
            device if isinstance(device, str) else str(device) if isinstance(device, Device) else str(jax.devices()[0])
        )
        self.jnp_array_kwargs = jnp_array_kwargs

    def _map_devices_to_str(self) -> Mapping[str, "jaxlib.xla_extension.Device"]:
        import jax

        return {str(device): device for device in jax.devices()}

    ...

    def _tensorize(self, value):
        ...

        with jax.default_device(self.device_mapping[self.device]):
            # calling jnp.array on a np.ndarray does copy the data
            # see https://github.com/google/jax/issues/4486
            return jnp.array(value, **{**default_dtype, **self.jnp_array_kwargs})
```

But note that the latter also throws a warning if the provided `device` is not a string but a `jaxlib.xla_extension.Device`, so that's why it needs to be converted to string.
2023-02-18T21:10:14Z
https://github.com/huggingface/datasets/pull/5547#issuecomment-1435769704
https://api.github.com/repos/huggingface/datasets/issues/5547
IC_kwDODunzps5VlBto
https://api.github.com/repos/huggingface/datasets/issues/comments/1435769704/reactions
2023-02-19T11:05:25Z
https://api.github.com/repos/huggingface/datasets/issues/comments/1435769704
https://avatars.githubusercontent.com/u/36760800?v=4
https://api.github.com/users/alvarobartt/events{/privacy}
https://api.github.com/users/alvarobartt/followers
https://api.github.com/users/alvarobartt/following{/other_user}
https://api.github.com/users/alvarobartt/gists{/gist_id}
https://github.com/alvarobartt
alvarobartt
MDQ6VXNlcjM2NzYwODAw
https://api.github.com/users/alvarobartt/orgs
https://api.github.com/users/alvarobartt/received_events
https://api.github.com/users/alvarobartt/repos
https://api.github.com/users/alvarobartt/starred{/owner}{/repo}
https://api.github.com/users/alvarobartt/subscriptions
User
https://api.github.com/users/alvarobartt
public
NONE
_The documentation is not available anymore as the PR was closed or merged._
2023-02-19T09:32:23Z
https://github.com/huggingface/datasets/pull/5547#issuecomment-1435937552
https://api.github.com/repos/huggingface/datasets/issues/5547
IC_kwDODunzps5VlqsQ
https://api.github.com/repos/huggingface/datasets/issues/comments/1435937552/reactions
2023-02-21T16:06:43Z
https://api.github.com/repos/huggingface/datasets/issues/comments/1435937552
https://avatars.githubusercontent.com/u/99929124?v=4
https://api.github.com/users/HuggingFaceDocBuilderDev/events{/privacy}
https://api.github.com/users/HuggingFaceDocBuilderDev/followers
https://api.github.com/users/HuggingFaceDocBuilderDev/following{/other_user}
https://api.github.com/users/HuggingFaceDocBuilderDev/gists{/gist_id}
https://github.com/HuggingFaceDocBuilderDev
HuggingFaceDocBuilderDev
U_kgDOBfTMJA
https://api.github.com/users/HuggingFaceDocBuilderDev/orgs
https://api.github.com/users/HuggingFaceDocBuilderDev/received_events
https://api.github.com/users/HuggingFaceDocBuilderDev/repos
https://api.github.com/users/HuggingFaceDocBuilderDev/starred{/owner}{/repo}
https://api.github.com/users/HuggingFaceDocBuilderDev/subscriptions
User
https://api.github.com/users/HuggingFaceDocBuilderDev
public
MEMBER
After some investigation, it seems that when using `device=jaxlib.xla_extension.Device` instead of `device=string` it shows the warning so that later formats fail as that cannot be unpickled.

So I think we can either add that specifically in `use_with_jax.mdx` documentation entry I'm creating at #5535 so that the users know that they need to surroung the `jaxlib.xla_extension.Device` with `str()`, or find a workaround to override default `deepcopy` behavior with `def __deepcopy__(self)` so that the device param is converted to string if provided as a `jaxlib.xla_extension.Device`, but not sure if the latter works 😕 

Do you think there's any other possible solution to this issue? Thanks, @lhoestq 
2023-02-19T11:09:03Z
https://github.com/huggingface/datasets/pull/5547#issuecomment-1435958981
https://api.github.com/repos/huggingface/datasets/issues/5547
IC_kwDODunzps5Vlv7F
https://api.github.com/repos/huggingface/datasets/issues/comments/1435958981/reactions
2023-02-19T11:09:03Z
https://api.github.com/repos/huggingface/datasets/issues/comments/1435958981
https://avatars.githubusercontent.com/u/36760800?v=4
https://api.github.com/users/alvarobartt/events{/privacy}
https://api.github.com/users/alvarobartt/followers
https://api.github.com/users/alvarobartt/following{/other_user}
https://api.github.com/users/alvarobartt/gists{/gist_id}
https://github.com/alvarobartt
alvarobartt
MDQ6VXNlcjM2NzYwODAw
https://api.github.com/users/alvarobartt/orgs
https://api.github.com/users/alvarobartt/received_events
https://api.github.com/users/alvarobartt/repos
https://api.github.com/users/alvarobartt/starred{/owner}{/repo}
https://api.github.com/users/alvarobartt/subscriptions
User
https://api.github.com/users/alvarobartt
public
MEMBER
Cool ! Specifying the device is indeed super important.


I think we can just require `device` to always be a string for now, and add an example in the doc on how to get the string that corresponds to a `jaxlib.xla_extension.Device` ? This way we never deal with objects that are not picklable
2023-02-19T14:03:11Z
https://github.com/huggingface/datasets/pull/5547#issuecomment-1435997610
https://api.github.com/repos/huggingface/datasets/issues/5547
IC_kwDODunzps5Vl5Wq
https://api.github.com/repos/huggingface/datasets/issues/comments/1435997610/reactions
2023-02-19T14:03:11Z
https://api.github.com/repos/huggingface/datasets/issues/comments/1435997610
https://avatars.githubusercontent.com/u/42851186?v=4
https://api.github.com/users/lhoestq/events{/privacy}
https://api.github.com/users/lhoestq/followers
https://api.github.com/users/lhoestq/following{/other_user}
https://api.github.com/users/lhoestq/gists{/gist_id}
https://github.com/lhoestq
lhoestq
MDQ6VXNlcjQyODUxMTg2
https://api.github.com/users/lhoestq/orgs
https://api.github.com/users/lhoestq/received_events
https://api.github.com/users/lhoestq/repos
https://api.github.com/users/lhoestq/starred{/owner}{/repo}
https://api.github.com/users/lhoestq/subscriptions
User
https://api.github.com/users/lhoestq
public
MEMBER
> Cool ! Specifying the device is indeed super important.
> 
> I think we can just require `device` to always be a string for now, and add an example in the doc on how to get the string that corresponds to a `jaxlib.xla_extension.Device` ? This way we never deal with objects that are not picklable

Sure, then I'll restrict it to string for now! Also regarding the documentation update, should we wait until #5535 is merged so that I add this on top of that?
2023-02-19T15:04:43Z
https://github.com/huggingface/datasets/pull/5547#issuecomment-1436012349
https://api.github.com/repos/huggingface/datasets/issues/5547
IC_kwDODunzps5Vl889
https://api.github.com/repos/huggingface/datasets/issues/comments/1436012349/reactions
2023-02-19T15:04:43Z
https://api.github.com/repos/huggingface/datasets/issues/comments/1436012349
https://avatars.githubusercontent.com/u/36760800?v=4
https://api.github.com/users/alvarobartt/events{/privacy}
https://api.github.com/users/alvarobartt/followers
https://api.github.com/users/alvarobartt/following{/other_user}
https://api.github.com/users/alvarobartt/gists{/gist_id}
https://github.com/alvarobartt
alvarobartt
MDQ6VXNlcjM2NzYwODAw
https://api.github.com/users/alvarobartt/orgs
https://api.github.com/users/alvarobartt/received_events
https://api.github.com/users/alvarobartt/repos
https://api.github.com/users/alvarobartt/starred{/owner}{/repo}
https://api.github.com/users/alvarobartt/subscriptions
User
https://api.github.com/users/alvarobartt
public
MEMBER
CI is failing due to missing `resampy` in `librosa` already being fixed by @lhoestq in https://github.com/huggingface/datasets/pull/5554
2023-02-20T20:17:01Z
https://github.com/huggingface/datasets/pull/5547#issuecomment-1437507712
https://api.github.com/repos/huggingface/datasets/issues/5547
IC_kwDODunzps5VrqCA
https://api.github.com/repos/huggingface/datasets/issues/comments/1437507712/reactions
2023-02-20T20:17:01Z
https://api.github.com/repos/huggingface/datasets/issues/comments/1437507712
https://avatars.githubusercontent.com/u/36760800?v=4
https://api.github.com/users/alvarobartt/events{/privacy}
https://api.github.com/users/alvarobartt/followers
https://api.github.com/users/alvarobartt/following{/other_user}
https://api.github.com/users/alvarobartt/gists{/gist_id}
https://github.com/alvarobartt
alvarobartt
MDQ6VXNlcjM2NzYwODAw
https://api.github.com/users/alvarobartt/orgs
https://api.github.com/users/alvarobartt/received_events
https://api.github.com/users/alvarobartt/repos
https://api.github.com/users/alvarobartt/starred{/owner}{/repo}
https://api.github.com/users/alvarobartt/subscriptions
User
https://api.github.com/users/alvarobartt
public
MEMBER
@lhoestq already moved to a global variable, I can confirm that the following now works:

```python
import copy
import pickle

import jax
import pyarrow as pa

from datasets.formatting import JaxFormatter


_COL_A = [0, 1, 2]
_COL_B = ["foo", "bar", "foobar"]
_COL_C = [[[1.0, 0.0, 0.0]] * 2, [[0.0, 1.0, 0.0]] * 2, [[0.0, 0.0, 1.0]] * 2]
pa_table = pa.Table.from_pydict({"a": _COL_A, "b": _COL_B, "c": _COL_C})

device = jax.devices()[0]
formatter = JaxFormatter(device=str(device))

pickle.dumps(formatter)
copy.deepcopy(formatter)
```
2023-02-21T14:59:41Z
https://github.com/huggingface/datasets/pull/5547#issuecomment-1438631426
https://api.github.com/repos/huggingface/datasets/issues/5547
IC_kwDODunzps5Vv8YC
https://api.github.com/repos/huggingface/datasets/issues/comments/1438631426/reactions
2023-02-21T14:59:41Z
https://api.github.com/repos/huggingface/datasets/issues/comments/1438631426
https://avatars.githubusercontent.com/u/36760800?v=4
https://api.github.com/users/alvarobartt/events{/privacy}
https://api.github.com/users/alvarobartt/followers
https://api.github.com/users/alvarobartt/following{/other_user}
https://api.github.com/users/alvarobartt/gists{/gist_id}
https://github.com/alvarobartt
alvarobartt
MDQ6VXNlcjM2NzYwODAw
https://api.github.com/users/alvarobartt/orgs
https://api.github.com/users/alvarobartt/received_events
https://api.github.com/users/alvarobartt/repos
https://api.github.com/users/alvarobartt/starred{/owner}{/repo}
https://api.github.com/users/alvarobartt/subscriptions
User
https://api.github.com/users/alvarobartt
public
MEMBER
> Looks all good now thank you !
> 
> Is there anything else you wanted to add ? Otherwise I think it's ready for merge

Nothing else to add, I've already applied your suggestions, so ready to merge! Thanks for your input/feedback @lhoestq :hugs:
2023-02-21T15:42:47Z
https://github.com/huggingface/datasets/pull/5547#issuecomment-1438699776
https://api.github.com/repos/huggingface/datasets/issues/5547
IC_kwDODunzps5VwNEA
https://api.github.com/repos/huggingface/datasets/issues/comments/1438699776/reactions
2023-02-21T15:43:02Z
https://api.github.com/repos/huggingface/datasets/issues/comments/1438699776
https://avatars.githubusercontent.com/u/36760800?v=4
https://api.github.com/users/alvarobartt/events{/privacy}
https://api.github.com/users/alvarobartt/followers
https://api.github.com/users/alvarobartt/following{/other_user}
https://api.github.com/users/alvarobartt/gists{/gist_id}
https://github.com/alvarobartt
alvarobartt
MDQ6VXNlcjM2NzYwODAw
https://api.github.com/users/alvarobartt/orgs
https://api.github.com/users/alvarobartt/received_events
https://api.github.com/users/alvarobartt/repos
https://api.github.com/users/alvarobartt/starred{/owner}{/repo}
https://api.github.com/users/alvarobartt/subscriptions
User
https://api.github.com/users/alvarobartt
public
NONE
<details>
<summary>Show benchmarks</summary>

PyArrow==6.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.009815 / 0.011353 (-0.001538) | 0.005443 / 0.011008 (-0.005565) | 0.101244 / 0.038508 (0.062736) | 0.036573 / 0.023109 (0.013464) | 0.304761 / 0.275898 (0.028863) | 0.365527 / 0.323480 (0.042047) | 0.008244 / 0.007986 (0.000258) | 0.004200 / 0.004328 (-0.000128) | 0.077471 / 0.004250 (0.073221) | 0.045266 / 0.037052 (0.008214) | 0.310213 / 0.258489 (0.051724) | 0.344247 / 0.293841 (0.050406) | 0.039530 / 0.128546 (-0.089016) | 0.012254 / 0.075646 (-0.063393) | 0.335039 / 0.419271 (-0.084233) | 0.049525 / 0.043533 (0.005992) | 0.298350 / 0.255139 (0.043211) | 0.312031 / 0.283200 (0.028832) | 0.108581 / 0.141683 (-0.033102) | 1.481178 / 1.452155 (0.029023) | 1.497662 / 1.492716 (0.004946) |

### 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.014762 / 0.018006 (-0.003244) | 0.447099 / 0.000490 (0.446609) | 0.009074 / 0.000200 (0.008874) | 0.000688 / 0.000054 (0.000633) |

### Benchmark: benchmark_indices_mapping.json

| metric | select | shard | shuffle | sort | train_test_split |
|--------|---|---|---|---|---|
| new / old (diff) | 0.027466 / 0.037411 (-0.009945) | 0.109715 / 0.014526 (0.095189) | 0.119062 / 0.176557 (-0.057495) | 0.188964 / 0.737135 (-0.548171) | 0.127057 / 0.296338 (-0.169282) |

### 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.395092 / 0.215209 (0.179883) | 3.948091 / 2.077655 (1.870436) | 1.795160 / 1.504120 (0.291040) | 1.603704 / 1.541195 (0.062509) | 1.714491 / 1.468490 (0.246001) | 0.700489 / 4.584777 (-3.884288) | 3.767493 / 3.745712 (0.021781) | 3.288374 / 5.269862 (-1.981488) | 1.783711 / 4.565676 (-2.781965) | 0.085119 / 0.424275 (-0.339156) | 0.012349 / 0.007607 (0.004742) | 0.502135 / 0.226044 (0.276091) | 5.019321 / 2.268929 (2.750392) | 2.236469 / 55.444624 (-53.208155) | 1.914376 / 6.876477 (-4.962101) | 1.998579 / 2.142072 (-0.143494) | 0.847841 / 4.805227 (-3.957386) | 0.166035 / 6.500664 (-6.334629) | 0.062469 / 0.075469 (-0.013000) |

### 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.245380 / 1.841788 (-0.596408) | 14.757872 / 8.074308 (6.683564) | 14.460373 / 10.191392 (4.268981) | 0.152981 / 0.680424 (-0.527443) | 0.029001 / 0.534201 (-0.505200) | 0.439597 / 0.579283 (-0.139686) | 0.437232 / 0.434364 (0.002868) | 0.532464 / 0.540337 (-0.007873) | 0.629225 / 1.386936 (-0.757711) |

</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.007165 / 0.011353 (-0.004188) | 0.005220 / 0.011008 (-0.005789) | 0.075849 / 0.038508 (0.037341) | 0.032717 / 0.023109 (0.009608) | 0.331205 / 0.275898 (0.055307) | 0.364955 / 0.323480 (0.041475) | 0.005518 / 0.007986 (-0.002468) | 0.004069 / 0.004328 (-0.000259) | 0.073900 / 0.004250 (0.069650) | 0.046346 / 0.037052 (0.009294) | 0.337473 / 0.258489 (0.078984) | 0.393062 / 0.293841 (0.099222) | 0.037533 / 0.128546 (-0.091013) | 0.012577 / 0.075646 (-0.063070) | 0.087975 / 0.419271 (-0.331297) | 0.049508 / 0.043533 (0.005975) | 0.333423 / 0.255139 (0.078284) | 0.354345 / 0.283200 (0.071145) | 0.099879 / 0.141683 (-0.041804) | 1.413304 / 1.452155 (-0.038851) | 1.494222 / 1.492716 (0.001506) |

### 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.206835 / 0.018006 (0.188828) | 0.438246 / 0.000490 (0.437757) | 0.000410 / 0.000200 (0.000210) | 0.000059 / 0.000054 (0.000004) |

### Benchmark: benchmark_indices_mapping.json

| metric | select | shard | shuffle | sort | train_test_split |
|--------|---|---|---|---|---|
| new / old (diff) | 0.028186 / 0.037411 (-0.009225) | 0.109322 / 0.014526 (0.094797) | 0.119581 / 0.176557 (-0.056975) | 0.191784 / 0.737135 (-0.545351) | 0.125100 / 0.296338 (-0.171238) |

### 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.419418 / 0.215209 (0.204209) | 4.167374 / 2.077655 (2.089720) | 1.995812 / 1.504120 (0.491693) | 1.804602 / 1.541195 (0.263407) | 1.869131 / 1.468490 (0.400641) | 0.709486 / 4.584777 (-3.875291) | 3.838019 / 3.745712 (0.092307) | 2.086206 / 5.269862 (-3.183656) | 1.323970 / 4.565676 (-3.241707) | 0.089477 / 0.424275 (-0.334798) | 0.012402 / 0.007607 (0.004795) | 0.519291 / 0.226044 (0.293246) | 5.194091 / 2.268929 (2.925162) | 2.487055 / 55.444624 (-52.957570) | 2.122495 / 6.876477 (-4.753982) | 2.194910 / 2.142072 (0.052837) | 0.842837 / 4.805227 (-3.962390) | 0.167229 / 6.500664 (-6.333435) | 0.064690 / 0.075469 (-0.010779) |

### 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.275931 / 1.841788 (-0.565857) | 14.577000 / 8.074308 (6.502692) | 13.633235 / 10.191392 (3.441843) | 0.184511 / 0.680424 (-0.495913) | 0.017439 / 0.534201 (-0.516762) | 0.424374 / 0.579283 (-0.154909) | 0.427803 / 0.434364 (-0.006561) | 0.527790 / 0.540337 (-0.012548) | 0.627301 / 1.386936 (-0.759635) |

</details>
</details>

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

2023-02-21T16:10:55Z
https://github.com/huggingface/datasets/pull/5547#issuecomment-1438743000
https://api.github.com/repos/huggingface/datasets/issues/5547
IC_kwDODunzps5VwXnY
https://api.github.com/repos/huggingface/datasets/issues/comments/1438743000/reactions
2023-02-21T16:10:55Z
https://api.github.com/repos/huggingface/datasets/issues/comments/1438743000
https://avatars.githubusercontent.com/in/15368?v=4
https://api.github.com/users/github-actions%5Bbot%5D/events{/privacy}
https://api.github.com/users/github-actions%5Bbot%5D/followers
https://api.github.com/users/github-actions%5Bbot%5D/following{/other_user}
https://api.github.com/users/github-actions%5Bbot%5D/gists{/gist_id}
https://github.com/apps/github-actions
github-actions[bot]
MDM6Qm90NDE4OTgyODI=
https://api.github.com/users/github-actions%5Bbot%5D/orgs
https://api.github.com/users/github-actions%5Bbot%5D/received_events
https://api.github.com/users/github-actions%5Bbot%5D/repos
https://api.github.com/users/github-actions%5Bbot%5D/starred{/owner}{/repo}
https://api.github.com/users/github-actions%5Bbot%5D/subscriptions
Bot
https://api.github.com/users/github-actions%5Bbot%5D
public
## What's in this PR?

After exploring for a while the JAX integration in 🤗`datasets`, I found out that, even though JAX prioritizes the TPU and GPU as the default device when available, the `JaxFormatter` doesn't let you specify the device where you want to place the `jax.Array`s in case you don't want to rely on JAX's default array placement.

So on, I've included the `device` param in `JaxFormatter` but there are some things to take into consideration:
* A formatted `Dataset` is copied with `copy.deepcopy` which means that if one adds the param `device` in `JaxFormatter` as a `jaxlib.xla_extension.Device`, it "fails" because that object cannot be serialized (instead of serializing the param adds a random hash instead). That's the reason why I added a function `_map_devices_to_str` to basically create a mapping of strings to `jaxlib.xla_extension.Device`s so that `self.device` is a string and not a `jaxlib.xla_extension.Device`.
* To create a `jax.Array` in a device you need to either create it in the default device and then move it to the desired device with `jax.device_put` or directly create it in the device you want with `jax.default_device()` context manager.
* JAX will create an array by default in `jax.devices()[0]`

More information on JAX device management is available at https://jax.readthedocs.io/en/latest/faq.html#controlling-data-and-computation-placement-on-devices

## What's missing in this PR?

I've tested it both locally in CPU (Mac M2 and Mac M1, as no GPU support for Mac yet), and in GPU and TPU in Google Colab, let me know if you want me to provide you the Notebook for the latter.

But I did not implement any integration test as I wanted to get your feedback first.