Fix RESOURCE_EXHAUSTED error when dealing with large datasets in Flax example scripts
Author: duongna21
Pull request: yes

# What does this PR do?

Fixes #15411. All Flax example scripts using `generate_batch_splits` raise RESOURCE_EXHAUSTED error when training with large datasets on TPU.

## Fix

Simply transfer the jnp array `samples_idx` back to the host before putting it in `np.split`.

## Who can review?
cc potential reviewers: @sgugger, @patil-suraj, @patrickvonplaten

Comments:
_The documentation is not available anymore as the PR was closed or merged._

@duongna21 Do you know the reason why we get `RESOURCE_EXHAUSTED` without the fix? Also, it would be nice to add a comment before the line to explain it a bit 🙏 

Also, does this issue occur only when we run on TPU? ~~I see this error before when running on TPU, but never see it on CPU.~~

Thank you!

@ydshieh I dug a little bit into the memory profiling when running the original code:

```
>>> import jax.numpy as jnp
>>> import numpy as np
>>> import time
>>> def generate_batch_splits(samples_idx: jnp.ndarray, batch_size: int) -> jnp.ndarray:
...     num_samples = len(samples_idx)
...     samples_to_remove = num_samples % batch_size
...     if samples_to_remove != 0:
...         samples_idx = samples_idx[:-samples_to_remove]
...     sections_split = num_samples // batch_size
...     batch_idx = np.split(samples_idx, sections_split)
...     jax.profiler.save_device_memory_profile("generate_batch_splits_tpu_bs64.prof")
...     return batch_idx
...  
>>> input_rng = jax.random.PRNGKey(0)
>>> samples_idx = jax.random.permutation(input_rng, jnp.arange(50758958)) # 50758958 is num of samples in my case
>>> tic = time.time()
>>> train_batch_idx = generate_batch_splits(samples_idx, batch_size=64) # per device batch size = 8
>>> toc = time.time()
>>> print('elapsed time: ', toc-tic)
elapsed time:  578.8218853473663
```

<img src="https://i.imgur.com/8tns1X2.png" width="400" height="300" />


=> **968.15MB** (keep increasing when we use larger dataset) of TPU memory ~~wasted~~ used while generating batch splits. When I transfered the jnp array `samples_idx` back to the CPU before putting it in `np.split`:
```
...
def generate_batch_splits(samples_idx: jnp.ndarray, batch_size: int) -> jnp.ndarray:
    samples_idx = jax.device_get(samples_idx)
    num_samples = len(samples_idx)
    ...
>>> print('elapsed time: ', toc-tic)
elapsed time:  3.2475574016571045
```
<img src="https://i.imgur.com/MCDDYHv.png" width="220" height="300" />

=> We just saved **968.15MB** (and 575 seconds). But we have an even better solution to save another **193.63MB** (and 2 seconds) by generating `samples_idx` using CPU in the first place:

```
...
samples_idx = np.random.permutation(np.arange(50758958))  # 50758958 is num of samples in my case
...
>>> print('elapsed time: ', toc-tic)
elapsed time:  1.0959625244140625
```

Anyway, 1GB of wasted TPU memory is not the decisive factor in causing RESOURCE_EXHAUTED, but this fix saves us from the weird RESOURCE_EXHAUTED error when using the same batch size and only increasing the dataset size.

Of course, this issue can occur on CPU when we run out of RAM.


Thanks a lot for the PR @duongna21 !
You're right we should better use numpy there - the PR looks good to me :-)