`MapDataset` does not respect `jax.default_device(jax.devices("cpu")[0])`.

[james77777778]

I would like to use CPU only for the transforms using jax.numpy.

However, I’ve found that MapDataset does not respect jax.default_device(jax.devices("cpu")[0]).

You can see a clear demonstration of this issue in the following Colab notebook:
https://colab.research.google.com/drive/1hRfWBe8a09bPNnXSVjeLvcrYjoqcUA2b?usp=sharing

Is there a reason for this behavior?
Could this be a bug?

[iindyk]

thanks for sharing the repro!

the main difference between the MapDataset and other scenarios is that your transform runs in background threads as opposed to the main thread. I modified the MapDataset scenario to disable this by adding to_iter_dataset(grain.Readoptions(prefetch_buffer_size=0)) and the device is CPU as expected.

do you know if jax.default_device is expected to be thread-local? this doesn't seem right

[james77777778]

the main difference between the MapDataset and other scenarios is that your transform runs in background threads as opposed to the main thread. I modified the MapDataset scenario to disable this by adding to_iter_dataset(grain.Readoptions(prefetch_buffer_size=0)) and the device is CPU as expected.

Thanks for the clarification!

do you know if jax.default_device is expected to be thread-local? this doesn't seem right

I believe jax.default_device is thread-local:
https://github.com/jax-ml/jax/blob/main/jax/_src/config.py#L1449-L1461

This might be the cause. Do you know how to force CPU computation in the transform or would it be better to use jax.device_put instead of jax.default_device?

Thanks!

[mhyatt000]

Hey I've been having this issue as well. More importantly, I cannot use mp_prefetch with jax operations since the subproc dont seem to respect jax.default_device(jax.devices("cpu")[0])

In my case I routinely hit CUDA OOM error since all the jax arrays move straight to the GPU unintentionally. This limits by loader to only 1 proc
Any feedback is welcome