Three classes, one instance check

Use isinstance(x, jax.Array) to find out whether you are holding real numbers and the answer will mislead you. Inside a jitted function the argument is a tracer, carrying a shape and a dtype and no values whatsoever, and the check still returns True. The class docstring promises exactly that, in a line written as an example: True both inside and outside traced functions.

So the check answers a narrower question than the one people ask it. It says the object speaks the array interface. It says nothing about whether anything has been computed, which is why it is useless as a guard against tracers and fine as a type annotation. Chapter 2 owns what a tracer is and what tracing discards; the fact that belongs here is only that a tracer passes this test.

A NumPy array fails the same check, and jnp functions accept one anyway. The union in the source that describes what they really take is ArrayLike, which is Array plus np.ndarray plus the Python and NumPy scalar types. What a function accepts and what isinstance admits are two different sets, and only the second one is jax.Array.

jnp.ndarray is not a second type either. It is the same object under a second name, so the NumPy-shaped spelling of the check is the identical check, down to object identity.

import jax
import jax.numpy as jnp
import numpy as np

x = jnp.arange(3.)
print(type(x).__name__, type(x).__mro__[1].__name__)
print(isinstance(x, jax.Array), isinstance(np.ones(3), jax.Array))
print(jnp.ndarray is jax.Array)

@jax.jit
def f(v):
    print(type(v).__name__, isinstance(v, jax.Array))
    return v * 2

f(x)

# ArrayImpl object
# True False
# True
# DynamicJaxprTracer True

jax.Array is an abstract base class with no implementation in it, and the classes that satisfy it arrive by two different routes. ArrayImpl, the concrete array jaxlib hands back, is registered as a virtual subclass from Python, which is why its method resolution order runs two entries long and lands on object. A TODO above the registration line says the plan is true inheritance at the C++ level, eventually.

Tracer takes the other route and inherits outright, so Array really does appear in its method resolution order. PRNGKeyArray, the type jax.random.key returns, inherits too. Chapter 8 owns what a key is and how to fold one; the fact for here is that it answers this instance check like anything else.

Three lines above the class statement, a comment says the type is not meant to include non-standard array types like KeyArray. The class statement in prng.py inherits from it directly. The comment predates that class and nothing updated it, so the class statements are the ones to read.

# jax/_src/basearray.py
# Array is a type annotation for standard JAX arrays and tracers produced by
# core functions in jax.lax and jax.numpy; it is not meant to include
# future non-standard array types like KeyArray and BInt.
class Array(abc.ABC):

# jax/_src/array.py, the last line of the module
# TODO(jakevdp) replace this with true inheritance at the C++ level.
basearray.Array.register(ArrayImpl)

# jax/_src/core.py
class Tracer(typing.Array, metaclass=StrictABCMeta):

# jax/_src/prng.py, the type the comment above says is not included
class PRNGKeyArray(jax.Array):

A concrete array is two things fastened together. An abstract value, the ShapedArray carrying shape and dtype, and a list of device buffers with one entry per addressable device. Ask a tracer for that buffer list and the attribute is not there at all, because a tracer is the first half and never the second.

Deleting proves the seam is real. x.delete() frees the buffers and leaves the Python object standing, so x.shape, x.dtype, and x.sharding all keep answering while the numbers are gone. Blocking on the same array afterward raises out of the runtime with a message naming a deleted or donated buffer. Donation, the other way buffers vanish under you, is chapter 11's to explain.

The metadata and the memory come apart. Deleting is the operation that takes them apart on purpose.

import jax
import jax.numpy as jnp

x = jnp.ones((4, 4))
print(x.is_deleted(), x.is_ready())
x.delete()
print(x.is_deleted())
print(x.shape, x.dtype, x.sharding)
try:
    x.block_until_ready()
except Exception as err:
    print(type(err).__name__)
    print(err)

# False True
# True
# (4, 4) float32 SingleDeviceSharding(device=CpuDevice(id=0), memory_kind=unpinned_host)
# XlaRuntimeError
# INVALID_ARGUMENT: BlockHostUntilReady() called on deleted or donated buffer

Two arrays can sit on the same device for two different reasons. One landed there because it had to land somewhere, the other because you said so, and the array remembers which. committed is a boolean the abstract base class requires of every implementation, and its docstring is the plainest statement of JAX's placement rules anywhere in the codebase.

jax.device_put(x, device) sets it. jax.device_put(x) with no device does not, and on an array already living on a device that call is the identity function, handing back the same object. Name a device, even the device the array is already sitting on, and a different object comes back with the bit set.

In that second case nothing about the bytes moved. What changed is a commitment: from here on, this array will not be relocated to meet another operand.

  @property
  @abc.abstractmethod
  def committed(self) -> bool:
    """Whether the array is committed or not.

    An array is committed when it is explicitly placed on device(s) via JAX
    APIs. For example, `jax.device_put(np.arange(8), jax.devices()[0])` is
    committed to device 0. While `jax.device_put(np.arange(8))` is uncommitted
    and will be placed on the default device.

    Computations involving some committed inputs will happen on the committed
    device(s) and the result will be committed on the same device(s).
    Invoking an operation on arguments that are committed to different device(s)
    will raise an error.
    """

Add an uncommitted array to a committed one and the result comes back committed, on the committed operand's device. That single rule is why one device_put near the top of a program can decide where a long chain of downstream arrays ends up, with none of the lines in between mentioning a device.

jax.default_device changes where new arrays land without committing any of them. Arrays built inside the context manager land on the device you named and stay uncommitted, so they remain free to be pulled somewhere else by the first committed operand they meet. Leave the block and the array keeps the device it got, still uncommitted.

Two uncommitted arrays on two different devices do not raise. They resolve onto the default device, quietly, and the result is uncommitted like both of its parents. Commitment is what turns a placement disagreement into an error, which the next section is about.

# run with: XLA_FLAGS=--xla_force_host_platform_device_count=2
import jax
import jax.numpy as jnp

d0, d1 = jax.devices()
u = jnp.ones(3)
print(u.committed, u.device)

c = jax.device_put(u, d1)
print(c.committed, c.device, c is u)
print(jax.device_put(u).committed)

with jax.default_device(d1):
    v = jnp.ones(3)
print(v.committed, v.device)

print((u + c).committed, (u + c).device)
print((u + v).committed, (u + v).device)

# False TFRT_CPU_0
# True TFRT_CPU_1 False
# False
# False TFRT_CPU_1
# True TFRT_CPU_1
# False TFRT_CPU_0

Commit two arrays to two devices, add them, and JAX refuses before anything runs. The message names the primitive, both shapes, and both device id lists, which is more detail than most placement problems give you. It is a plain ValueError, not a runtime error out of the device like the one the deleted buffer produced.

Read what the message leaves out. It never uses the word committed, so the sentence you get is a report of where the arguments were, not a diagnosis of why they could not move. The diagnosis is the previous section: each argument had been pinned, and a pinned array does not migrate to meet another operand.

Which makes the fix mechanical once you can name the cause. Either drop one of the commitments, or add a device_put that moves one argument onto the other's device on purpose, and the copy becomes a line you wrote instead of a line JAX refused to write for you.

# run with: XLA_FLAGS=--xla_force_host_platform_device_count=2
import jax
import jax.numpy as jnp

d0, d1 = jax.devices()
a = jax.device_put(jnp.ones(3), d0)
b = jax.device_put(jnp.ones(3), d1)
try:
    a + b
except Exception as err:
    print(type(err).__name__)
    print(err)

# ValueError
# Received incompatible devices for jitted computation. Got argument x of add
# with shape float32[3] and device ids [0] on platform CPU and argument y of
# add with shape float32[3] and device ids [1] on platform CPU