assert jnp.ndarray trigger host to device copy on GPU/TPU?

[luweizheng]

Hi all,

I am on GPU platform. I am wondering how to do assert with JAX and prevent host device data copy?

If assert return True, then it will continue the following graph computation. If assert return False, it will raise exceptions.

i = jnp.asarray(1)
steps_num = jnp.asarray(3)

# will it do memory copy from device to host?
assert i < steps_num

# if is a Python control flow, I guess it will trigger memory copy from deivce to host.
if i < steps_num:
     print("True")

Which way should I use to prevent memory copy and do assert ?

[jakevdp]

When you do assert xyz or if xyz, Python will call the __bool__ method of the object xyz.

In the case of two Device Arrays x and y, the statement z = x < y will return another Device Array of type bool. The __bool__ method for DeviceArrays is defined here: https://github.com/google/jax/blob/b3a62cd3f2be15a7ed23771b371835e2977961be/jax/_src/device_array.py#L269

This currently will forward z._value.__bool__, and z._value involves converting the buffer to a py buffer on the CPU: https://github.com/google/jax/blob/b3a62cd3f2be15a7ed23771b371835e2977961be/jax/_src/device_array.py#L144-L150

I don't think there's any way around this: Python control flow happens on the CPU, because Python runs on the CPU, so you cannot assert the value of something that doesn't live on the CPU. But this is a single byte being transferred in this case and only at trace-time (not at run-time), so I suspect it's not worth worrying about.

[luweizheng]

When you do assert xyz or if xyz, Python will call the __bool__ method of the object xyz.

In the case of two Device Arrays x and y, the statement z = x < y will return another Device Array of type bool. The __bool__ method for DeviceArrays is defined here:

https://github.com/google/jax/blob/b3a62cd3f2be15a7ed23771b371835e2977961be/jax/_src/device_array.py#L269

This currently will forward z._value.__bool__, and z._value involves converting the buffer to a py buffer on the CPU:

https://github.com/google/jax/blob/b3a62cd3f2be15a7ed23771b371835e2977961be/jax/_src/device_array.py#L144-L150

I don't think there's any way around this: Python control flow happens on the CPU, because Python runs on the CPU, so you cannot assert the value of something that doesn't live on the CPU. But this is a single byte being transferred in this case and only at trace-time (not at run-time), so I suspect it's not worth worrying about.

Thanks! Just find tf.assert_greater method. I guess this method is designed to prevent host device copy?

[jakevdp]

I don't think tf.assert_greater will work with JAX.

If you can say more about what exactly you're trying to accomplish, we might be able to help find a solution.