ShardedDeviceArray from tf.data.Dataset.shard

[krzysztofrusek]

Hi, I have a question regarding the data pipeline in multi-host training.
In particular, I have multiple workers equipped with GPU and each worker can access central data strorage.
I would like to use tf.data.Dataset.shard to load part of the batch independently on each worker and join the shard in a single ShardedDeviceArray that can be handled by pmap.

It looks like jax.device_put_sharded does the job but it requires a list of shards on the host, and I want to load them independently on the workers.

I imagine my loop to be like

for it in ds.shard(n,jax.process_index()):
  sharded_batch = unknow_function(it,...)
  pmaped_train_step(sharded_batch)

What is the most efficient way to do it?