Understanding data commitment

[ayaka14732]

In the JAX documentation (Controlling data and computation placement on devices), the word "committed" is mentioned. It is used to describe whether a JAX array has been "committed" to a device.

For a worked-out example, I have read through test_computation_follows_data in multi_device_test.py.

My questions:

1. What does the word "committed" exactly mean here?
2. More specifically, I am trying to perform some computations on a TPU device while loading new data onto that device. Is it possible to perform the two operations concurrently in JAX?

[yashk2810]

What does the word "committed" exactly mean here?

committed means if a jax.Array is pinned to that device. Is there any particular line in that test which is confusing?

Is it possible to perform the two operations concurrently in JAX?

Yeah it should be possible since JAX's dispatch is asynchronous. To be accurate, can you provide an example?

[ayaka14732]

committed means if a jax.Array is pinned to that device. Is there any particular line in that test which is confusing?

I couldn't quite understand. If a jax.Array is not pinned to that device, where does it reside? Besides, how is it different from block_until_ready?

[yashk2810]

It will reside on the default device which currently is the 0th local device. If a jax.Array is uncommitted, Jax is free to move it to other devices but if it is committed, then Jax won't move it.

For example if x is committed on d0 and y is committed on d1 and you do x + y, then jax will error. But if x is uncommitted on d0, the jax will move it to d1 to perform the addition. In this case the result will be on d1 committed since y was committed.

I don't understand the .block_until_ready() question 😅

[remifan]

Thanks @yashk2810 for the clarification. It is what "the computation follows data placement" mean in the FAQ.

[ayaka14732]

I have realized that I had mistakenly thought that committed indicated whether a data transfer had completed, which is why I associated it with block_until_ready(). I now understand that an uncommitted array still resides on the default device, but its placement is not fixed and it may follow a committed input, while block_until_ready() only waits for asynchronous work to complete. This answers my question. Thank you!