>>> import jax, jax.numpy as jnp
>>> x = jnp.asarray(0)
>>> def f(x): return jnp.zeros(0, device=x.device)
>>> f(x)
Array([], shape=(0,), dtype=float32)
>>> jax.jit(f)(x)
AttributeError: DynamicJaxprTracer has no attribute device
>>> def g(x): return jnp.zeros(0, device=next(iter(x.devices())))
>>> jax.jit(g)(x)
ConcretizationTypeError: Abstract tracer value encountered where concrete value is expected: traced array with shape int32[]
The devices() method was called on traced array with shape int32[].
The error occurred while tracing the function g at <ipython-input-6-5e654833a7fe>:1 for jit. This concrete value was not available in Python because it depends on the value of the argument x.
The error occurred while tracing the function g at <ipython-input-6-5e654833a7fe>:1 for jit. This concrete value was not available in Python because it depends on the value of the argument x.

def f(x):
    xp = array_namespace(x)
    return xp.asarray(123, device=x.device)