IterableDataset with_format does not support 'device' keyword for jax
COLLABORATOR
Hi! Yes, only `torch` is currently supported. Unlike `Dataset`, `IterableDataset` is not PyArrow-backed, so we cannot simply call `to_numpy` on the underlying subtables to format them numerically. Instead, we must manually convert examples to (numeric) arrays while preserving consistency with `Dataset`, which is not trivial, so this is still a to-do.
2023-03-16T18:22:48Z
https://github.com/huggingface/datasets/issues/5637#issuecomment-1472508425
https://api.github.com/repos/huggingface/datasets/issues/5637
IC_kwDODunzps5XxLIJ
https://api.github.com/repos/huggingface/datasets/issues/comments/1472508425/reactions
2023-03-16T18:22:48Z
https://api.github.com/repos/huggingface/datasets/issues/comments/1472508425
https://avatars.githubusercontent.com/u/47462742?v=4
https://api.github.com/users/mariosasko/events{/privacy}
https://api.github.com/users/mariosasko/followers
https://api.github.com/users/mariosasko/following{/other_user}
https://api.github.com/users/mariosasko/gists{/gist_id}
https://github.com/mariosasko
mariosasko
MDQ6VXNlcjQ3NDYyNzQy
https://api.github.com/users/mariosasko/orgs
https://api.github.com/users/mariosasko/received_events
https://api.github.com/users/mariosasko/repos
https://api.github.com/users/mariosasko/starred{/owner}{/repo}
https://api.github.com/users/mariosasko/subscriptions
User
https://api.github.com/users/mariosasko
public
NONE
Any plans to support it in the future? Or would streaming dataset be left without support for jax and tensorflow?
2023-03-16T18:30:59Z
https://github.com/huggingface/datasets/issues/5637#issuecomment-1472528579
https://api.github.com/repos/huggingface/datasets/issues/5637
IC_kwDODunzps5XxQDD
https://api.github.com/repos/huggingface/datasets/issues/comments/1472528579/reactions
2023-03-16T18:30:59Z
https://api.github.com/repos/huggingface/datasets/issues/comments/1472528579
https://avatars.githubusercontent.com/u/91322985?v=4
https://api.github.com/users/Lime-Cakes/events{/privacy}
https://api.github.com/users/Lime-Cakes/followers
https://api.github.com/users/Lime-Cakes/following{/other_user}
https://api.github.com/users/Lime-Cakes/gists{/gist_id}
https://github.com/Lime-Cakes
Lime-Cakes
MDQ6VXNlcjkxMzIyOTg1
https://api.github.com/users/Lime-Cakes/orgs
https://api.github.com/users/Lime-Cakes/received_events
https://api.github.com/users/Lime-Cakes/repos
https://api.github.com/users/Lime-Cakes/starred{/owner}{/repo}
https://api.github.com/users/Lime-Cakes/subscriptions
User
https://api.github.com/users/Lime-Cakes
public
NONE
@mariosasko 
`IterableDataset` and `with_format('torch', device=device)` also does not work (using `datasets==3.2.0`)
2025-01-07T06:59:20Z
https://github.com/huggingface/datasets/issues/5637#issuecomment-2574519430
https://api.github.com/repos/huggingface/datasets/issues/5637
IC_kwDODunzps6ZdAyG
https://api.github.com/repos/huggingface/datasets/issues/comments/2574519430/reactions
2025-01-07T06:59:33Z
https://api.github.com/repos/huggingface/datasets/issues/comments/2574519430
https://avatars.githubusercontent.com/u/113588658?v=4
https://api.github.com/users/amitportnoy/events{/privacy}
https://api.github.com/users/amitportnoy/followers
https://api.github.com/users/amitportnoy/following{/other_user}
https://api.github.com/users/amitportnoy/gists{/gist_id}
https://github.com/amitportnoy
amitportnoy
U_kgDOBsU5sg
https://api.github.com/users/amitportnoy/orgs
https://api.github.com/users/amitportnoy/received_events
https://api.github.com/users/amitportnoy/repos
https://api.github.com/users/amitportnoy/starred{/owner}{/repo}
https://api.github.com/users/amitportnoy/subscriptions
User
https://api.github.com/users/amitportnoy
public
### Describe the bug

As seen here: https://huggingface.co/docs/datasets/use_with_jax dataset.with_format() supports the keyword 'device', to put data on a specific device when loaded as jax. However, when called on an IterableDataset, I got the error `TypeError: with_format() got an unexpected keyword argument 'device'`

Looking over the code, it seems IterableDataset support only pytorch and no support for jax device keyword? 
https://github.com/huggingface/datasets/blob/fc5c84f36684343bff3e424cb0fd1ac5ecdd66da/src/datasets/iterable_dataset.py#L1029



### Steps to reproduce the bug

1. Load an IterableDataset (tested in streaming mode)
2. Call with_format('jax',device=device)

### Expected behavior

I expect to call `with_format('jax', device=device)`  as per [documentation](https://huggingface.co/docs/datasets/use_with_jax) without error

### Environment info

Tested with installing newest (dev) and also pip release (2.10.1).
- `datasets` version: 2.10.2.dev0
- Platform: Linux-5.15.89+-x86_64-with-debian-bullseye-sid
- Python version: 3.7.12
- Huggingface_hub version: 0.12.1
- PyArrow version: 11.0.0
- Pandas version: 1.3.5
