Introduction

Errors around TensorFlow stratified sampling usually come from one of two problems: either the labels and target distribution do not match the API's expectations, or the code is using an older queue-based pattern in a newer TensorFlow environment. The fix depends on which world your code belongs to. Legacy TensorFlow 1 style code has strict shape and session requirements, while modern TensorFlow 2 projects are usually better served by the tf.data resampling tools instead of an older stratified_sample call.

Understand What Stratified Sampling Needs

Stratified sampling means you want to draw examples by class according to some desired distribution. For that to work, TensorFlow needs three things to line up:

• a feature tensor or dataset element
• a label tensor that clearly identifies the class
• a target distribution that matches the number of classes

If labels are missing, have the wrong dtype, or contain values outside the expected class range, the operation often fails with shape or validation errors.

A Modern TensorFlow 2 Approach

In TensorFlow 2, the cleanest solution is usually to use rejection_resample on a Dataset. This avoids many of the old queue-runner issues.

1import tensorflow as tf
2
3features = tf.constant([[1.0], [2.0], [3.0], [4.0], [5.0], [6.0]])
4labels = tf.constant([0, 0, 0, 0, 1, 1], dtype=tf.int64)
5
6dataset = tf.data.Dataset.from_tensor_slices((features, labels))
7
8resampled = dataset.rejection_resample(
9    class_func=lambda x, y: y,
10    target_dist=[0.5, 0.5],
11    seed=42,
12)
13
14for sampled_class, (x, y) in resampled.take(4):
15    print(sampled_class.numpy(), x.numpy(), y.numpy())

This example resamples the dataset toward a balanced class distribution. The key detail is that target_dist must align with the classes returned by class_func. If you have three label values, the target distribution must describe three classes.

Why Older stratified_sample Code Often Breaks

Many examples on old forums use queue-based TensorFlow 1 code. That style assumes graph execution, sessions, queue runners, and label tensors with very specific shapes and types. If you paste that code into a TensorFlow 2 notebook with eager execution enabled, it often breaks for reasons that have nothing to do with the sampling math.

If you truly need compatibility mode, keep the entire pipeline in the TensorFlow 1 style.

1import tensorflow as tf
2
3tf.compat.v1.disable_eager_execution()
4
5labels = tf.constant([0, 0, 1, 1], dtype=tf.int32)
6values = tf.constant([[10.0], [20.0], [30.0], [40.0]], dtype=tf.float32)
7
8dataset = tf.data.Dataset.from_tensor_slices((values, labels)).repeat(1)
9iterator = tf.compat.v1.data.make_one_shot_iterator(dataset)
10next_item = iterator.get_next()
11
12with tf.compat.v1.Session() as sess:
13    try:
14        while True:
15            print(sess.run(next_item))
16    except tf.errors.OutOfRangeError:
17        pass

This example does not call the old stratified_sample API directly, but it shows the environment you need if your code still depends on session-based TensorFlow patterns. Mixing compatibility-mode graph code with eager-style dataset iteration is a common source of confusion.

What to Check When You See an Error

Start with labels. They should usually be integer class IDs such as 0, 1, and 2, not one-hot encoded vectors unless the API explicitly expects that form. Then inspect the target distribution. Its length must match the number of classes you expect to sample.

Next, verify that feature tensors and label tensors have the same leading dimension. A mismatch between the number of examples and the number of labels will often show up as a shape error, even though the root cause is simple misalignment.

Finally, check whether the code example you copied was written for TensorFlow 1 or TensorFlow 2. That single detail explains many stratified_sample failures.

Common Pitfalls

One common mistake is passing floating-point labels or string labels into an API that expects integer class IDs. Convert the labels first so TensorFlow can map them cleanly to class buckets.

Another frequent mistake is building a target distribution that does not sum to a sensible whole or does not match the class count. If you have labels 0 and 1, a two-element distribution is required. If you add a third class later, the distribution must change with it.

Developers also often copy deprecated TensorFlow 1 queue-based code into TensorFlow 2 without disabling eager execution or without replacing the old API with tf.data resampling tools. In that situation, even valid sampling logic fails because the runtime model is wrong.

Finally, remember that resampling changes the data distribution seen by the model. That can help with imbalance, but it also changes training dynamics. Validate the model after the fix instead of assuming a balanced sampler automatically improves results.

Summary

• Most TensorFlow stratified sampling errors come from label mismatches, wrong target distributions, or mixing TensorFlow 1 and TensorFlow 2 execution styles.
• In TensorFlow 2, tf.data resampling tools are usually a better choice than older queue-based patterns.
• Labels should be in a clear class-ID form and aligned with the feature tensor length.
• The target distribution must match the number of classes being sampled.
• If legacy code is unavoidable, keep the whole pipeline in tf.compat.v1 session mode instead of mixing execution models.