Ensuring that a large matrix stays in CPU memory

[SNMS95]

Hi,

I have a function where a list of JAX arrays needs to be vertically stacked. The problem is that if I use jax.numpy.vstack to do the stacking, I run out of GPU memory. On the other hand, if I use ordinary numpy, the stacking takes place in CPU (where I have enough memory) but causes tracer error if I try to do any transformation.

I tried wrapping the onp.vstack with a pure_callback plus a custom_jvp rule. However, the result of the forward pass seems to end up in GPU memory anyway.

Finally, I tried to put the list of arrays into cpu using device_put() and then use jax.numpy.vstack. But I am not sure if that large matrix would always stay in CPU if I perform transformations on the function (like a reverse-mode autodiff).

def func(a, cpu):
  inp = [a*np.ones((2, 2)), a*np.ones((2, 2))]
  inp_cpu = jax.device_put(inp, cpu)
  val = jax.device_put(np.vstack(inp_cpu), cpu)
  return val

a = 2.
cpu = jax.devices("cpu")[0]
#func(a, cpu).device_buffer.device() #---------> This returns the object in CPU
#jax.jacfwd(func)(a, cpu).device_buffer.device() #---------> This also returns the object in CPU
jax.jacrev(func)(a, cpu).device_buffer.device() # -----> This gives the result in GPU

Is there a way to ensure that a particular computation will always take place in CPU, while being JAX transformable?

[jakevdp]

You may be able to address this by specifying jax.default_device:

with jax.default_device(cpu):
  result = jax.jacrev(func)(a, cpu)