Introduction

Dataset.shard is part of TensorFlow’s tf.data API and is used to split one dataset stream into multiple non-overlapping pieces. It is especially useful in distributed training, where each worker should process only its own subset of the data instead of reading the entire dataset.

The core idea is simple: if you have num_shards workers, each worker takes one shard identified by its index. TensorFlow then keeps every num_shards-th element that belongs to that shard.

What shard Actually Does

The method signature is:

dataset = dataset.shard(num_shards, index)

This does not create contiguous chunks by file or by range. Instead, it selects dataset elements by position.

For example:

1import tensorflow as tf
2
3dataset = tf.data.Dataset.range(10)
4shard = dataset.shard(num_shards=3, index=1)
5
6print(list(shard.as_numpy_iterator()))

Output:

[1, 4, 7]

TensorFlow keeps elements whose position satisfies the shard pattern for index=1 among three shards.

Why This Is Useful

Suppose you have four workers in a distributed job. Without sharding, each worker may read the same full dataset and duplicate training work unnecessarily. With shard(4, worker_id), each worker sees only its assigned subset.

That helps with:

• distributed training efficiency
• avoiding duplicate data consumption
• deterministic worker partitioning
• reducing unnecessary input-pipeline work

Sharding is also useful outside distributed training if you want to split a dataset consistently across processes or experimental runs.

Basic Multiworker Pattern

A common manual pattern looks like this:

1import tensorflow as tf
2
3
4def make_dataset(worker_index: int, num_workers: int):
5    dataset = tf.data.Dataset.range(20)
6    dataset = dataset.shard(num_shards=num_workers, index=worker_index)
7    dataset = dataset.batch(4)
8    return dataset
9
10
11for batch in make_dataset(worker_index=2, num_workers=4):
12    print(batch.numpy())

Each worker passes its own worker_index, and all workers use the same num_workers.

Place shard Early in the Pipeline

In many pipelines, sharding should happen early, before expensive mapping or randomization. That way each worker only preprocesses the data it actually needs.

1import tensorflow as tf
2
3
4dataset = tf.data.Dataset.range(1000)
5dataset = dataset.shard(num_shards=8, index=3)
6dataset = dataset.map(lambda x: x * x, num_parallel_calls=tf.data.AUTOTUNE)
7dataset = dataset.batch(32)
8dataset = dataset.prefetch(tf.data.AUTOTUNE)

This is usually better than mapping first and sharding later, because later sharding can waste compute on items another worker will discard.

Interaction with Shuffle

Ordering matters. If you shard before shuffle, each worker shuffles only its own subset. If you shuffle first and then shard, you may still get correct partitioning, but the workers may each read more of the upstream pipeline than necessary.

A common practical rule is:

• shard early for efficiency
• shuffle after sharding if each worker should randomize its own data stream

That keeps the pipeline cheaper while still giving good training behavior in many setups.

Using shard with File Datasets

shard works on dataset elements regardless of whether those elements are numbers, file paths, or parsed records.

1import tensorflow as tf
2
3files = tf.data.Dataset.from_tensor_slices([
4    "file1.tfrecord",
5    "file2.tfrecord",
6    "file3.tfrecord",
7    "file4.tfrecord",
8])
9
10worker_files = files.shard(num_shards=2, index=0)
11print(list(worker_files.as_numpy_iterator()))

This is a useful way to distribute file processing across workers before opening the files.

shard Versus Automatic Sharding

TensorFlow distributed strategies can also apply automatic sharding in some training configurations. Manual Dataset.shard is still valuable when you want explicit control over exactly how input partitioning happens.

That said, do not combine manual and automatic sharding blindly. If both are active without intention, you may end up with unexpectedly small or incomplete data slices.

Common Pitfalls

One common mistake is expecting shard to return one contiguous chunk of the dataset. It actually selects elements by position pattern, not by one continuous range.

Another issue is applying shard too late in the pipeline, after expensive preprocessing. That wastes resources because each worker processes data that may later be thrown away.

It is also easy to misconfigure num_shards and index. The index must be in the range from 0 to num_shards - 1.

Finally, be careful when using manual sharding together with distributed training features that may already perform automatic input sharding.

Summary

• 'Dataset.shard(num_shards, index) keeps every num_shards-th dataset element for the given shard index.'
• It is especially useful when multiple workers should process disjoint subsets of the same dataset.
• Sharding early in the pipeline is usually more efficient than sharding after expensive transformations.
• It works for numeric data, records, and file-path datasets alike.
• Manual sharding gives explicit control, but it should be coordinated with any automatic sharding behavior in distributed training.