Introduction

When you have several tf.data.Dataset objects and want to draw examples from them in random order, the right tool is usually tf.data.Dataset.sample_from_datasets. It mixes multiple input streams probabilistically, which is much better than simple concatenation when you need balanced sampling or controlled proportions.

Use sample_from_datasets for Probabilistic Mixing

The core pattern is straightforward:

1import tensorflow as tf
2
3dataset_a = tf.data.Dataset.from_tensor_slices(["a0", "a1", "a2"]).repeat()
4dataset_b = tf.data.Dataset.from_tensor_slices(["b0", "b1", "b2"]).repeat()
5
6mixed = tf.data.Dataset.sample_from_datasets(
7    [dataset_a, dataset_b],
8    weights=[0.7, 0.3],
9)
10
11for item in mixed.take(10):
12    print(item.numpy().decode("utf-8"))

This creates a new dataset that draws from dataset_a about 70 percent of the time and from dataset_b about 30 percent of the time.

That is the usual answer when the goal is "randomly sample from multiple datasets."

Why Not Just Concatenate and Shuffle?

Concatenating datasets and then shuffling is fine when all examples can be materialized together and you only care about random order. It is not the same as weighted sampling, especially when:

• the datasets have different sizes
• you want oversampling of rare classes
• the datasets are infinite or repeated streams
• you want a stable probability mix over time

sample_from_datasets operates at the dataset level, not only at the final merged element list.

A Common Class-Balancing Use Case

Suppose one dataset contains positive examples and another contains negative examples. If the raw data is imbalanced, you can use equal sampling weights to present a more balanced stream to the model.

1positive = tf.data.Dataset.from_tensor_slices([1, 1, 1]).repeat()
2negative = tf.data.Dataset.from_tensor_slices([0, 0, 0, 0, 0, 0]).repeat()
3
4balanced = tf.data.Dataset.sample_from_datasets(
5    [positive, negative],
6    weights=[0.5, 0.5],
7)
8
9print(list(balanced.take(8).as_numpy_iterator()))

Even though the negative dataset is larger, the resulting training stream can stay balanced because the sampling probabilities are controlled explicitly.

Repeat Behavior Matters

If one dataset is finite and the others repeat forever, the mixed pipeline may stop earlier than you expect unless you manage repetition carefully. In training pipelines, it is common to call .repeat() on all input datasets before sampling so that the combined stream does not end after one source is exhausted.

You should also think about whether exhaustion should stop sampling or whether other datasets should continue. The desired behavior depends on whether the inputs represent:

• training sources with replacement
• a fixed evaluation corpus
• class-specific streams that must stay aligned

Add Local Shuffle Before Sampling

Sampling chooses which dataset to draw from, but it does not automatically randomize the order inside each input stream. If each source has meaningful ordering, shuffle the sources too:

1dataset_a = dataset_a.shuffle(1000).repeat()
2dataset_b = dataset_b.shuffle(1000).repeat()
3
4mixed = tf.data.Dataset.sample_from_datasets(
5    [dataset_a, dataset_b],
6    weights=[0.5, 0.5],
7).batch(32).prefetch(tf.data.AUTOTUNE)

That gives you randomness at both levels:

• random source selection
• random element order inside each source

interleave Solves a Different Problem

interleave is often mentioned in the same discussions, but it serves a different purpose. It is good for reading from many files or input streams concurrently with structured cycling behavior. If you want probabilistic source selection, sample_from_datasets is usually the clearer API.

Common Pitfalls

• Concatenating and shuffling when you actually need weighted per-dataset sampling.
• Forgetting to repeat() datasets in training and then wondering why the mixed pipeline ends early.
• Assuming dataset-level sampling also shuffles item order inside each dataset.
• Using weights that do not reflect the actual training objective, which can create accidental bias.
• Reaching for interleave when the requirement is random source choice rather than parallel stream reading.

Summary

• Use tf.data.Dataset.sample_from_datasets to randomly draw from multiple datasets.
• Weights let you control the long-run sampling proportion from each source.
• Shuffle individual datasets if their internal order matters.
• Repeat datasets deliberately in training pipelines so the mixed stream behaves as intended.
• 'interleave is useful for concurrent reading, but it is not the same as probabilistic dataset sampling.'