Is there a way to force jax into CPU-only mode?

[Jacob-Stevens-Haas]

At the top of my file I have

import jax
jax.config.update("jax_default_device", "cpu")

Yet somehow my computation stalls with a warning that the GPU is out of memory:

2025-05-07 10:22:53.156746: W external/xla/xla/tsl/framework/bfc_allocator.cc:501] Allocator (GPU_0_bfc) ran out of memory trying to allocate 335.67MiB (rounded to 351973632)requested by op 
2025-05-07 10:22:53.156791: W external/xla/xla/tsl/framework/bfc_allocator.cc:512] *___________________________________________________________________________________________________
E0507 10:22:53.156803 4025642 pjrt_stream_executor_client.cc:2839] Execution of replica 0 failed: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 351973512 bytes. [tf-allocator-allocation-error='']

I verified that jax.config.__getattribute__("jax_default_device") is indeed CPU. The problem was because a context manager set it as GPU temporarily. I'm wondering whether there's a way to specifically prohibit jax from seeing a particular device. I'm also wondering what would happen if the GPU was truly out of memory so that jax couldn't even allocate the 1% of GPU it uses even when the default device is CPU.

[selamw1]

By default, JAX places new arrays on jax.devices()[0], which is often the first GPU  or TPU.

To make JAX ignore GPUs entirely, you could try setting the environment variable JAX_PLATFORMS="cpu". This needs to be done before JAX initializes. Here are two common ways:

1. From your shell:

$ JAX_PLATFORMS=cpu python -c "import jax"

2. In your Python script (before the first import jax, or after restarting your Python kernel if jax was already imported):

import os
os.environ['JAX_PLATFORMS'] = 'cpu'
print(os.environ.get('JAX_PLATFORMS')) # cpu

import jax
print(jax.devices()) # [CpuDevice(id=0)]

Regarding your question about a GPU being truly out of memory: if JAX (without JAX_PLATFORMS="cpu") tries to initialize on such a GPU, its GPU backend would likely fail to start. JAX would then operate as if no GPU were present. However, using JAX_PLATFORMS="cpu" preempts this entirely by telling JAX not to even look for GPUs.