Is there a clean way to update a sharded array such that each shard gets a portion of the update

[sash-a]

Say I have a pre-allocated and sharded array x that I want to live on TPU cores. Then on the CPU I generate an array y (len(y) < len(x)) and want to update x with the values in y.

Assuming some curr_idx for x and shapes of x=(16,2) and y=(4,2) I could simply do:

x = x.at[curr_idx:curr_idx+len(y),:].set(y)

This correctly places y on device and in the sharded array, but each shard does not get an equal slice of y. If I had 2 TPU cores, how could I place y in x such that each shard of x gets a (2,2) slice of y?

P.S: I see that a sharded array has the global_shards attribute and I could get the slice indexes from there. But this means I would likely have to do multiple device_put calls (one for each shard). Just wondering if there is a cleaner way to do this?