Generating a array of random choices on the CPU, not GPU

[Peter-Vincent]

I want to generate an array of choice, using the random.choice function, but I would like to place that array on the CPU, rather than the GPU (which is my current default backend).

As I understand it, I can place arrays, once they are created, on my device of choice (for example the CPU) with the following code (see here)

device_cpu = devices('cpu')[0]
new_arr = jax.device_put(arr, device=device_cpu)

Or I can force a function to be executed on the CPU like this (see here)

cpu_func = jax.jit(func, device=device_cpu)

The problem is that the function I want to use is jax.random.choice.

Option 1 above is not viable, since just the creation of the device on the GPU exhausts the GPU memory, so I can't create it first, then move it to where I want. Option 2 doesn't work, since when jitting the function I encounter a jax.errors.ConcretizationTypeError, which makes total sense (the key error is printed below)

jax.errors.ConcretizationTypeError: Abstract tracer value encountered where concrete value is expected: traced array with shape int32[].
The error occurred in jax.random.choice()

I also can't see an option for the inputs to jax.random.choice for devive, that would let me place the array on the CPU by default, like I have seen for some other functions.

One final detail, I'm not passing arrays as arguments to the function, I'm only passing integers.

How can I call this function so that it creates the array of choices on the CPU rather than the GPU, when my default backend is GPU?

[Peter-Vincent]

Update -

My current solution is to initialise the array in batches, as follows

for i in range(num_batches):
    choice_list.append(jax.device_put(random.choice(subkeys[i],max_val,shape),device=device_cpu))
choices = jnp.concatenate(choice_list,axis=1)

And this results in the overall large array being placed on the CPU. But, there feels like there should be a more elegant solution to this problem?

[dfm]

The reason for that ConcretizationTypeError is that some of the arguments to random.choice must be static. You should be able to get your proposed jit version to work using something like:

cpu_func = jax.jit(random.choice, static_argnums=(1, 2), device=device_cpu)
cpu_func(key, max_val, shape)

Hope this helps!

[Peter-Vincent]

Oh yeah, thats a good idea. Thanks!