Can we use vmap without taking data to GPU?

[mohamad-amin]

Hey!

I have a function that I'm applying vmap on it.

Inside the vmap, I have to use jax.numpy as it is apparently dealing with TraceShapedArrays and at some point I'm using scipy.linalg.cho_solve which is being called using TraceShapedArrays as inputs and it can not accept them. Thus, I have to use jax.scipy.linalg.cho_solve but using this automatically takes the cho_solve computation to GPU, even though the input to it is not a jax.numpy array. This is problematic as my data that gets fed into cho_solve can not fit in any GPU. Thus, I would like to use vmap but without operating in GPU. Is there anyway to do this?

P.S: I need GPUs for some pmap computations before this vmap that doesn't fit on GPU, so I can't force jax to use CPU devices all the time.

Sincerely,
Amin