Introduction

Filtering a tf.data.Dataset is a normal part of building an input pipeline. The important detail is that the predicate passed to Dataset.filter must operate on tensors and return a scalar boolean tensor, not a regular Python bool.

Filter With Tensor Operations

The basic pattern is to define a function, or lambda, that receives one dataset element and returns True or False as a TensorFlow value.

1import tensorflow as tf
2
3values = tf.data.Dataset.from_tensor_slices(tf.constant([0, 1, 2, 3, 4, 5]))
4
5filtered = values.filter(lambda x: tf.equal(x % 2, 0))
6
7for item in filtered.as_numpy_iterator():
8    print(item)

This prints only the even values. The key call is tf.equal(x % 2, 0). That expression yields a tensor, so it can be compiled into the input pipeline.

If you write lambda x: x % 2 == 0, TensorFlow may still convert it correctly in simple cases, but using explicit TensorFlow ops is the safer habit because it keeps the predicate graph-friendly.

Filter Structured Elements

Real datasets often yield tuples such as feature and label pairs. In that case, your predicate must accept the same structure produced by the dataset.

1import tensorflow as tf
2
3features = tf.constant([[1.0], [2.0], [3.0], [4.0]])
4labels = tf.constant([0, 1, 0, 1])
5
6ds = tf.data.Dataset.from_tensor_slices((features, labels))
7
8only_positive_class = ds.filter(lambda x, y: tf.equal(y, 1))
9
10for x, y in only_positive_class:
11    print(x.numpy(), y.numpy())

This keeps only rows whose label equals 1. Notice that the feature tensor x is still available to the predicate even though the condition is based only on y.

The same idea works for dictionary-shaped records. If each element is a mapping with keys such as "id" and "label", filter by reading the appropriate field and returning a tensor condition.

Filter By A Set Of Allowed Values

A frequent requirement is to keep only elements whose ID belongs to a small allowlist. For a tiny list of values, tf.reduce_any with tf.equal is fine.

1import tensorflow as tf
2
3ids = tf.data.Dataset.from_tensor_slices(tf.constant([10, 11, 12, 13, 14]))
4allowed = tf.constant([11, 14])
5
6filtered = ids.filter(
7    lambda x: tf.reduce_any(tf.equal(x, allowed))
8)
9
10for item in filtered.as_numpy_iterator():
11    print(item)

For larger lookups, use a hash table so the membership check stays readable and efficient.

1import tensorflow as tf
2
3keys = tf.constant([101, 205, 309], dtype=tf.int32)
4values = tf.constant([True, True, True])
5initializer = tf.lookup.KeyValueTensorInitializer(keys, values)
6table = tf.lookup.StaticHashTable(initializer, default_value=False)
7
8ds = tf.data.Dataset.from_tensor_slices(tf.constant([100, 101, 102, 205, 400]))
9filtered = ds.filter(lambda x: table.lookup(x))
10
11for item in filtered.as_numpy_iterator():
12    print(item)

That pattern is useful when filtering examples by user ID, category ID, or any other discrete key.

Put Filtering In The Right Place

Filtering early in the pipeline usually saves work. If you shuffle, decode, augment, or batch elements before filtering, the pipeline does unnecessary processing on rows that will be dropped later.

A practical order is often:

1. read raw examples
2. parse them
3. filter unwanted rows
4. map expensive transformations
5. batch and prefetch

There are exceptions. Sometimes you need to parse a serialized record before the predicate can inspect its label or metadata. But once the necessary fields exist, filter as soon as possible.

Debugging A Filter Predicate

If a filter keeps nothing, inspect a few elements before the filter and verify dtypes. String, integer, and floating-point comparisons can fail silently when the expected type does not match the actual one.

A useful debugging trick is to temporarily convert the predicate into a map that returns the original value plus the boolean condition. That lets you inspect what the condition is computing before you drop elements.

Also remember that Dataset.filter expects one scalar boolean per element. Returning a boolean vector, or a Python container of booleans, will raise an error.

Common Pitfalls

The biggest pitfall is using regular Python control flow inside the predicate. Expressions such as if x in my_list or x in [1, 2, 3] operate in Python space and do not translate cleanly into TensorFlow graph execution.

Another common mistake is filtering after batching when the intent was to remove individual examples. Once the dataset is batched, the predicate receives a batch tensor instead of one record, so the logic must change.

People also get tripped up by dtype mismatches. Comparing an int64 tensor against an int32 constant can lead to confusing errors or implicit casts you did not intend.

Finally, do not overuse tf.py_function for filtering. It can work, but it makes the pipeline harder to optimize, serialize, and debug.

Summary

• 'Dataset.filter requires a predicate that returns a scalar boolean tensor.'
• Use TensorFlow ops such as tf.equal, tf.reduce_any, and lookup tables.
• Match the predicate arguments to the structure of each dataset element.
• Filter as early as practical to avoid wasted work downstream.
• Check dtypes and element structure first when the filter behaves unexpectedly.