Executing with GPU as the default device affects subsequent executions with CPU as the default device. How to debug?

[muupan]

EDIT2: I asked about the issue in numpyro's issues too: pyro-ppl/numpyro#1496.

EDIT3: The numpyro dependency has been removed from the code so that the issue is reproducible with jax (and numpy) only.

I have a function f that uses jax and numpyro (https://github.com/pyro-ppl/numpyro). I have a problem where execution with GPU as the default device affects subsequent executions with CPU as the default device. This only happens when JIT enabled.

I could not share the content of f, but it must be stateless. Here is the issue:

# cpu
with jax.default_device(jax.devices("cpu")[0]):
    y1 = f()
# cpu again
with jax.default_device(jax.devices("cpu")[0]):
    y2 = f()
# gpu
with jax.default_device(jax.devices("gpu")[0]):
    y3 = f()
# cpu yet again
with jax.default_device(jax.devices("cpu")[0]):
    y4 = f()

Here, I expect y1 == y2 == y4 since they are executed on the same device. However, while y1 == y2, y4 is different from them, implying that the in-gpu computation of y3 affects the in-cpu computation of y4. Indeed, even if I comment out y1 and y2, it does not affect y4. If comment out y3, y1 == y2 == y4 If I set JAX_DISABLE_JIT=1, then the issue disappears, with y1 == y2 == y4, which makes me suspect that this is related to JIT cache.

Since I haven't shared the content of f, I guess it is difficult to see what is going on inside it just from the information above, but I'd appreciate any advices or hints on what is going on or how to debug it.

I'm using jax==0.3.21 and jaxlib==0.3.20+cuda11.cudnn8.

EDIT: Here is the code to reproduce the issue. I'm using numpyro==0.10.1.

import jax
import jax.example_libraries.optimizers
import jax.numpy as jnp
import numpy as np


def loss_fn(params, data):
    t_train, n = data.shape
    mu = params["mu"]

    def transition(carry, t):
        h_prev, u_prev = carry
        h_t = 0.1 * (u_prev ** 2) + h_prev
        scale = jnp.sqrt(h_t)
        y_t = data[t]
        u_t = y_t - mu
        return (h_t, u_t), (mu, scale, y_t)

    h_0 = jnp.ones(n) * 0.01
    u_0 = jnp.ones(n)
    carry_init = (h_0, u_0)
    _, (mu, scale, y_t) = jax.lax.scan(transition, carry_init, jnp.arange(1, t_train))
    normalize_term = jnp.log(jnp.sqrt(2 * jnp.pi) * scale)
    value_scaled = (y_t - mu) / scale
    return (0.5 * value_scaled ** 2 + normalize_term).sum()


def f():
    np.random.seed(0)
    optimizer = jax.example_libraries.optimizers.adam(1.0)
    data = np.random.normal(loc=0, scale=0.01, size=(20, 1))
    data_jnp = jnp.array(data)
    params = {"mu": jnp.ones(data_jnp.shape[-1])}
    opt_state = optimizer[0](params)

    @jax.jit
    def step(i, opt_state):
        params = optimizer[2](opt_state)
        loss_value, grads = jax.value_and_grad(loss_fn)(params, data_jnp)
        opt_state = optimizer[1](i, grads, opt_state)
        return opt_state, loss_value

    for i in range(1000):
        opt_state, loss_value = step(i, opt_state)
        # if i % 100 == 0:
        #     print(f"step {i}, loss: {loss_value}")

    params = optimizer[2](opt_state)

    return params


with jax.default_device(jax.devices("cpu")[0]):
    y1 = f()

with jax.default_device(jax.devices("cpu")[0]):
    y2 = f()

with jax.default_device(jax.devices("gpu")[0]):
    y3 = f()

with jax.default_device(jax.devices("cpu")[0]):
    y4 = f()

print(f"{y1=}")
print(f"{y2=}")
print(f"{y3=}")
print(f"{y4=}")

output (without JAX_DISABLE_JIT=1):

y1={'mu': DeviceArray([0.05775195], dtype=float32)}
y2={'mu': DeviceArray([0.05775195], dtype=float32)}
y3={'mu': DeviceArray([0.045999], dtype=float32)}
y4={'mu': DeviceArray([-0.03762827], dtype=float32)}

output (with JAX_DISABLE_JIT=1):

y1={'mu': DeviceArray([-0.00767463], dtype=float32)}
y2={'mu': DeviceArray([-0.00767463], dtype=float32)}
y3={'mu': DeviceArray([-0.04210557], dtype=float32)}
y4={'mu': DeviceArray([-0.00767463], dtype=float32)}