jit does not respect backend type when the function has no arguments

[romanngg]

Example:

from jax.api import jit
import jax.random as random

def x():
  return random.normal(random.PRNGKey(1), (2, 3))

x().device_buffer.device()
jit(x)().device_buffer.device()
jit(x, backend='cpu')().device_buffer.device()

All these output GpuDevice(id=0), while the last one should be CPU.

[jekbradbury]

Functions without arguments can't meaningfully be jitted, since jit relies on data dependence to trace primitives inside the function (it injects tracer values for each function argument, and follows them through the function). That is, the return value of x() is actually a constant that was computed op-by-op (on the default device) when jit attempted to trace x(), and not the result of a jitted XLA computation at all.

We don't yet have a way to specify the device used for op-by-op computations; it might make sense for such a mechanism to be dynamically scoped (i.e., a context manager) since we can't use data dependence for scoping the way we do for jit.

[levskaya]

If for some reason you really need to use this idiom, you can force data dependence on constants (sometimes handy to prevent them from being folded out) by using lax.tie_in:

from jax.api import jit
import jax.random as random
from jax import lax

dummy = 0

def x(_placeholder):
  return lax.tie_in(_placeholder, random.normal(random.PRNGKey(1), (2, 3))) 

print(x(dummy).device_buffer.device())  # --> gpu
print(jit(x)(dummy).device_buffer.device())  # --> gpu
print(jit(x, backend='cpu')(dummy).device_buffer.device())  # --> cpu