google-deepmind / google-deepmind/acme
Avoid inlining large arrays in JaxInMemoryRandomSampleIterator
Open
- Dominant language
- Python
- Stars
- 4.1k
- Forks
- 553
- PR merge metrics
- No merged PRs in 30d
Description
The JaxInMemoryRandomSampleIterator currently inlines the in-memory dataset. See
https://github.com/deepmind/acme/blob/master/acme/datasets/tfds.py#L199-L200
This causes some OOM issues due to some issues in XLA and also when running on GPU the process might hang. I have filed a more detailed issue in the JAX project https://github.com/google/jax/issues/14080 and the authors recommend not inlining the array instead. I can create a PR if the developers would like to fix that.
Contributor guide
Assessment
This issue has not been assessed yet.