Specify CPU backend after JAX loaded

[cool-RR]

Is there a way to tell JAX to only use the CPU and not the GPU? I've previously used JAX_PLATFORMS=cpu but this appears to only work when you set it before JAX was initialized. I want a solution that works after JAX was imported.

[selamw1]

Even after JAX has been imported, you can still force it to use the CPU using the following methods:

1. Set a Default Device for JAX Operations

• You can explicitly set the default backend to the CPU by updating JAX’s configuration:

import jax

jax.config.update("jax_default_device", jax.devices("cpu")[0])

# Create an array and perform a computation on the CPU
x = jnp.ones((3, 3))
y = jnp.linalg.inv(x + jnp.eye(3))  
print(y.device) 
# TFRT_CPU_0

2. Manually Specify CPU for Computation

• If you want to execute specific computations on the CPU while keeping GPU active for other tasks, you can explicitly place tensors and computations on the CPU:

import jax
import jax.numpy as jnp

cpu_device = jax.devices("cpu")[0]

# Explicitly place an array on the CPU
x = jax.device_put(jnp.array([1, 2, 3]), cpu_device)

# Run computation on CPU
with jax.default_device(cpu_device):
    y = jnp.dot(x, x)

print(y.device)
# TFRT_CPU_0

[cool-RR]

Wonderful, thank you!