How to gather data from certain devices

[asparius]

me = jax.lax.axis_index("batch")
loss = 0
loss += my_loss(my_data)
for device_id in range(len(jax.device_count())):
     skip = jnp.not_equal(device_id,me)
     other_data = gather(my_data,device_id) # gather function works properly
     loss += skip * other_loss(my_data, other_data)

How can I modify this to gather data only from devices with a higher index than the current device? Following is a failing attempt, and by the way this code is pmapped and jitted.

me = jax.lax.axis_index("batch")
loss = 0
loss += my_loss(my_data)
for device_id in range(me.astype(int),len(jax.device_count())):
     skip = jnp.not_equal(device_id,me)
     other_data = gather(my_data,device_id) # gather function works properly
     loss += skip * other_loss(my_data, other_data)

If that's possible, is there any practical efficiency gain in doing that? Theoretically, it has the same complexity even though we only do half of it. Thanks in advance :)

[asparius]

Up!