Cannot change `jax_platform_name` on the fly

I am using jax.config.update('jax_platform_name', 'cpu') to change the device placement.

import jax
import jax.numpy as np
a = np.zeros((128, 128))
print(a.device())  # TPU
jax.config.update('jax_platform_name', 'cpu')
a = np.zeros((128, 128))
print(a.device())  # TPU, expected CPU

Related: #8879

[ayaka14732]

@theovincent

The correct way is:

import jax
import jax.numpy as jnp

print(jnp.ones(3).device())

with jax.default_device(jax.devices('cpu')[0]):
    print(jnp.ones(3).device()) 

[hawkinsp]

jax_platform_name isn’t intended to be set at run time. I suspect what you actually want to control is the default device placement, for which @skye was adding a context-manager-like API.

(I don’t believe it landed yet, though.)