When are JAX arrays committed to a device?

[mjo22]

I am struggling to understand how JAX manages when arrays are committed to a particular sharding. I have a very simple question about this example: https://docs.jax.dev/en/latest/notebooks/host-offloading.html#data-placement-with-device-put.

When the call to jnp.arange is made, was this data committed to a default device, or is somehow the computation not materialized until the sharding is specified? Also, which sharding did the underlying computation? Was this done on the default device then transferred to both host and device shardings, or something else?

In general, my question is twofold: what is happening under the hood before a sharding is specified? and when a sharding is never specified, when does JAX commit to a device? I’d like to understand how to minimize unnecessary memory transfers.

I am trying to understand both jax.device_put and how the in_sharding argument works for JAX array inputs to jax.jit.

[jakevdp]

When you create an array (in eager mode) it will be materialized on the default device, but will not be "committed", which means that when you operate on it, the compiler is free to move it to whichever device or sharding is most efficient. You can see that here:

import jax
jax.config.update("jax_num_cpu_devices", 2)

x = jax.numpy.arange(10)
print(f"{x.device=}\n{x.committed=}\n")

x.device=CpuDevice(id=0)
x.committed=False

When you later call device_put on that array, the buffer will be materialized on the specified device, and will also be considered "committed" to that device, meaning that the compiler will respect that data placement for any operations performed on the array:

x2 = jax.device_put(x, jax.devices()[1])
print(f"{x2.device=}\n{x2.committed=}\n")

x2.device=CpuDevice(id=1)
x2.committed=True

Note that the above is true outside JIT (i.e. when executing eagerly). Within JIT, when you create an array using jax.numpy.arange(10), the buffer of values may or may not ever actually be materialized on the device, depending on how the compiler chooses to optimize the program.

Does that help answer your question?

[mjo22]

This is very helpful. I assume it works the same way for jnp.asarray casts in eager vs under JIT?

[jakevdp]

jnp.asarray is an interesting case. It does several different things depending on the context:

• if it is passed a single non-JAX object (e.g. a NumPy array or a Python int), it will call device_put on them and move them, uncommitted, to the default device.
• if it is passed a single JAX array, it will return it unmodified (unless you pass dtype, ndmin or the like, in which case it will call reshape or convert_element_type as appropriate)
• if it is passed a list or nested list of objects, it will first call asarray on each individual entry, then call lax.concatenate one or more times to stack them into a single array of the appropriate shape.

Under JIT, the input will almost always already be a (traced) JAX array, in which case it is basically a no-op unless reshape, convert_element_type, or concatenate are necessary.