Missing `.device` attribute inside `@jax.jit`

[crusaderky]

Description

JAX arrays are missing the .device attribute when running inside jax.jit, and accessing the .devices() method raises a concretization error:

>>> 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.

This breaks Array API compatibility and hinders array API compliant libraries that use the pattern

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

You can see two such use cases in array-api-extra: https://github.com/search?q=repo%3Adata-apis%2Farray-api-extra+device%3D_compat.device&type=code

Workaround

I'm provisionally implementing a workaround in array-api-compat that causes device(x) to return None and to_device(x, device) to accept None to work around this issue. This will however produce outputs on the wrong device when x is not on the default device.

System info (python version, jaxlib version, accelerator, etc.)

JAX 0.4.35

[crusaderky]

It looks like jax.jit is propagating the device from the input arrays when they don't lie on the default device, but only if the input arrays are part of the graph that produces the output arrays:

>>> a = jnp.asarray([1,1,1,2,2])
>>> a.device
CudaDevice(id=0)
>>> def f(x): return jnp.zeros(0)
>>> def g(x): return jnp.where(x > 0, jnp.zeros(1), jnp.ones(1))
>>> cpu = jax.devices("cpu")[0]
>>> jax.jit(f)(a.to_device(cpu)).device
CudaDevice(id=0)
>>> jax.jit(g)(a.to_device(cpu)).device
CpuDevice(id=0)

[jakevdp]

Thanks for raising this: this is working as intended; there's no reasonable value to return for arr.device when arr is traced. arr in this context represents a logical entity that may or may not exist after the compilation pass, and even if it does exist it may not have any meaningful device! Given that, we thought the most reasonable implementation would be to make arr.device an AttributeError.

I hadn't appreciated before that this makes traced JAX code incompatible with the array API, but that makes sense now that you raise the issue.

It looks like jax.jit is propagating the device from the input arrays when they don't lie on the default device

This device propagation is a property of the function being compiled, not a property of the individual tracer used to define that function.

It sounds like if we want JAX code to be compatible with the array API when used with transformations like jit and vmap, we need tracer.device to be implemented and to return some value indicating that the concept of "device" is meaningless in the given context. @yashk2810 what do you think?

[crusaderky]

Can't we make the device of the input arrays static, along with their dtype and shape?

[crusaderky]

It sounds like if we want JAX code to be compatible with the array API when used with transformations like jit and vmap, we need tracer.device to be implemented and to return some value indicating that the concept of "device" is meaningless in the given context. @yashk2810 what do you think?

This is basically what I hacked in data-apis/array-api-compat#238.
The problem with it is that it causes the device of the output arrays to be unpredictable, whereas most users would expect it to match the input's.

The pattern affected is this:

@jax.jit
def f(x):
    xp = array_namespace(x)
    if some_fast_exit_condition:
        return xp.zeros(1, dtype=x.dtype, device=x.device)  # Or asarray, or ones, etc.
    else:
        return g(x)

See https://github.com/search?q=repo%3Adata-apis%2Farray-api-extra+device%3D_compat.device&type=code for two distinct examples of this pattern in action.

With my current workaround, if x is not on the default device then the output of f(x) will be sometimes on x.device and some other times on the default device.

[jakevdp]

x.device may make sense in the simplest of cases where x corresponds to an input array, but in general arr.device for a traced value arr is not something that is sensible to query. In the most extreme case, the buffer represented by arr may not even exist in the compiled computation (e.g. if the operation is part of a compiler fusion): how do we return a device for a buffer that doesn't exist?

[yashk2810]

As Jake said, we don't have device inside tracing context because it leads to all sorts of issues with caches and generally, it's against tracing concept since internally, JAX tries to abstractify everything.

But there is an idea of having AbstractDevice sort of concept though I don't see how it can help here.

[crusaderky]

But there is an idea of having AbstractDevice sort of concept though I don't see how it can help here.

It would help in not having to maintain a hack in array-api-compat in the long run.
It would not help in respecting the user's expectation to propagate the device from input to output when creating new objects.

[pearu]

It would help in not having to maintain a hack in array-api-compat in the long run.

In the case of JAX (as well as of other libraries that support JIT compilation for accelerator devices), Python Array API standard is being enforced on pure Python functions and on jitted functions, roughly speaking.
As mentioned above, enforcing device support (as specified by the Array API std) within jitted functions is impossible/inapplicable for inevitable reasons. Hence, in general, enforcing Python Array API standard on jitted functions cannot be 100 % not even in the long run.

So, I suggest that the issue can be resolved by documenting the level of JAX conformance to Python Array API standard for pure Python functions and jitted functions separately.

^ @rgommers

[rgommers]

I'm not sure I understand all the comments. .device isn't needed to evaluate to a concrete device object like CudaDevice(id=0), but to express the semantics of device handling. For example, "create a new array of ones with given shape on the same device as that of the input array" is expressed like:

def func(x1):
    x2 = jnp.ones((3, 2), device=x1.device)
    # now use x2 .... (can be combined with x1 efficiently because we know they're both on the same device)

In the most extreme case, the buffer represented by arr may not even exist in the compiled computation (e.g. if the operation is part of a compiler fusion): how do we return a device for a buffer that doesn't exist?

This is an example where I don't quite see the problem. The buffer being fused away is a pure performance optimization that does not affect the semantics of the original code.

[pearu]

@rgommers , your example does not have a problem. This issue is about the following example

@jax.jit
def func(x1):
    x2 = jnp.ones((3, 2), device=x1.device)

where x1.device is not defined inside the definition of jitted function.

Within a jitted function, "create a new array of ones with given shape on the same device as that of the input array" is pointless in the sense that the opposite "create a new array of ones with given shape on the another device as that of the input array" is not supported (it would require graph break support that would trigger device memory management tasks outside of the kernel execution).

[jakevdp]

The buffer being fused away is a pure performance optimization that does not affect the semantics of the original code.

Maybe so, but it does make it quite difficult to answer the question of which device the non-existent buffer is stored on!

I suspect the best resolution here would be for tracer.device to return some sentinel that represents an undefined device; and then when passed to a routine that creates a new tracer, that new tracer would also be on an undefined device (as all tracers are).

[crusaderky]

Maybe so, but it does make it quite difficult to answer the question of which device the non-existent buffer is stored on!

Does it? There are only four possible answers AFAIK:

1. the device of one of the input args of the jitted function, or
2. the default device, or
3. a device queried out of jax.devices, or
4. a device passed as a static parameter to the jitted function or from a global variable.

I expect n.3 and n.4 to be of very little practical interest so I suggest to disregard them.
To support n. 1 and 2, you'd need to trace the provenance of the future device through the graph.

[jakevdp]

I agree in the simplest cases you could apply the device from an input array to the created array. But were we to implement this, we'd have to worry about how to implement every case, not just the simplest case.

you'd need to trace the provenance of the future device through the graph.

This is difficult to do, because device placement is ultimately decided by the compiler, and the tracing pass (where we choose what arr.device returns) happens before the compilation pass. It certainly would be possible to build some abstract provenance graph during tracing in order to pass that to the compiler so that the compiler can populate the devices of intermediate arrays, but that's not the model JAX currently uses, so implementing it would require a major change to JAX's tracing and compilation architecture, and we'd need to show that the added complexity is justified.

And there's still the question I keep bringing up: what if the compiler decides to elide a certain operation, then the buffer (and therefore its device placement) do not exist? I'm not sure what to do in that case, even if we had some new abstract provenance graph available.