Run tests with CPU mode when on CPU+GPU version

I have built the package with

python build/build.py --enable_march_native --enable_mkl_dnn --enable_cuda
pip install -e build
pip install -e .

I think this should build a CPU + GPU version.

When I tried to run pytest -n 2 tests examples -W ignore, some of the tests failed. It appeared to me that some of the functions are only support in CPU mode.

I also tried JAX_ENABLE_X64=1 JAX_NUM_GENERATED_CASES=100 pytest -n auto tests and JAX_ENABLE_X64=0 JAX_NUM_GENERATED_CASES=100 pytest -n auto tests, but neither of them worked.

[skye]

Setting JAX_PLATFORM_NAME is the best way to do this. You can also manually device_put values and/or use the backend argument to various functions like pmap to run things on CPU, but the env var is the best way to set the platform globally (e.g. for performance comparisons).

[matthewfeickert]

When both CPU and GPU are available and no jax_platform_name flag is set, we preferentially select the GPU backend. After we merge #690, you’ll be able to write this to set the backend from the shell (if you have a posix shell):

JAX_PLATFORM_NAME=cpu pytest -n auto tests -W ignore

@mattjj So does this mean that if you want to be able to toggle using the GPU and the CPU (use case: do a performance comparison of the speedup you get with code using JAX on the GPU (without having to switch between virtual environments with two versions installed)) the preferred way of doing this is by setting the environment value? So either setting it through the Shell like you show above, or by doing

import os
os.environ["JAX_PLATFORM_NAME"] = "cpu"

? Or is there recommended API?