Missing sharding support in TENSORFLOW mode with Datasets?

[winston-zillow]

In the mnist_dist.py example, we have

ds = tf.data.TextLineDataset(files).map(parse_fn).batch(args.batch_size)

There is no shard(num_workers, worker_index) before the batching. I think that would mean all workers will be reading the same training data?

Though in order to invoke shard() the node needs to know the total number of workers through the context passing to the map function.

(I can provide a PR, though I want to see if this is indeed an issue.)

[leewyang]

For reference, this issue refers to this code.

Yes, for simplicity, this code just demonstrates each of the workers training on a dataset created from the same source files. As you indicated, it should be easy to shard by worker_index. Feel free to provide a PR. (Note: for a more complete example, it should probably include some sort of mechanism for shuffling the data per epoch... but that might be a little trickier, depending on implementation).

[leewyang]

Added sharding, shuffling, and epoch support in #254