How to set default device as CPU

[zohimchandani]

I would like to benchmark the performance of my code on CPU and GPU but Jax natively recognizes the GPU. How can I force it to run the code on CPU?

[YouJiacheng]

Option 1

import jax
import os
os.environ['CUDA_VISIBLE_DEVICES'] = ''
# before execute any computation / allocation
print(jax.numpy.ones(3).device()) # TFRT_CPU_0

Option 2

import jax
jax.config.update('jax_platform_name', 'cpu')
# before execute any computation / allocation
print(jax.numpy.ones(3).device()) # TFRT_CPU_0

Option 3

import jax

def f():
    return jax.numpy.ones(3)

print(jax.jit(f, backend='cpu')().device()) # TFRT_CPU_0

[zohimchandani]

Trying option 2:

[jakevdp]

Note that this setting will only work if it is run before any other JAX code in the session – the device cannot be changed via configuration flags mid-session. If you're using Jupyter, try restarting your kernel and then rerun the code.

[zohimchandani]

@jakevdp
Perfect, got it, any ideas what TFRT stands for?

Also, given that my remote machine has 20 CPUs, can pmap be used to parallelize the compuation over the CPUs or does it only support GPUs?

[YouJiacheng]

I think it is https://github.com/tensorflow/runtime/blob/master/documents/tfrt_host_runtime_design.md.
There is a config option to split the TFRT_CPU_0, which already represent all CPU cores, into multiple device for debugging purpose. But the performance is inferior.

os.environ['XLA_FLAGS'] = '--xla_force_host_platform_device_count=8'