TensorFlow
tf.data.Dataset
data preprocessing
variable length lists
machine learning

How to input a list of lists with different sizes in tf.data.Dataset

Master System Design with Codemia

Enhance your system design skills with over 120 practice problems, detailed solutions, and hands-on exercises.

Introduction

tf.data.Dataset is excellent for input pipelines, but many real datasets contain sequences with different lengths. Text token lists, click streams, and event histories rarely fit into a neat rectangular tensor. To use them correctly, you need to represent the varying length explicitly and choose a batching strategy that matches your model.

Why Uneven Nested Lists Need Special Handling

A regular TensorFlow tensor is dense, which means every row must have the same number of elements. If you try to build a tensor from Python lists with different inner lengths, TensorFlow cannot infer a single rectangular shape.

python
1import tensorflow as tf
2
3samples = [
4    [1, 2, 3],
5    [4, 5],
6    [6, 7, 8, 9],
7]
8
9try:
10    dense = tf.constant(samples)
11    print(dense)
12except Exception as exc:
13    print(type(exc).__name__)
14    print(exc)

That failure is expected. tf.data.Dataset.from_tensor_slices works best when it can slice a well-formed tensor, so irregular nested lists usually need a different entry point.

Build the Dataset With from_generator

The most flexible approach is tf.data.Dataset.from_generator. Each example is yielded one at a time, and the output signature declares that the sequence length is unknown.

python
1import tensorflow as tf
2
3samples = [
4    [1, 2, 3],
5    [4, 5],
6    [6, 7, 8, 9],
7]
8
9def gen():
10    for row in samples:
11        yield row
12
13ds = tf.data.Dataset.from_generator(
14    gen,
15    output_signature=tf.TensorSpec(shape=(None,), dtype=tf.int32),
16)
17
18for row in ds:
19    print(row.numpy())

The important detail is shape=(None,). That tells TensorFlow each row is a one-dimensional tensor whose length may vary.

Batch Variable-Length Rows With Padding

Many models still expect batched dense tensors. In that case, the usual solution is padded_batch, which pads each sequence in the batch to the batch-local maximum length.

python
1import tensorflow as tf
2
3samples = [
4    [1, 2, 3],
5    [4, 5],
6    [6, 7, 8, 9],
7]
8
9def gen():
10    for row in samples:
11        yield row
12
13ds = tf.data.Dataset.from_generator(
14    gen,
15    output_signature=tf.TensorSpec(shape=(None,), dtype=tf.int32),
16)
17
18batched = ds.padded_batch(
19    batch_size=2,
20    padded_shapes=(None,),
21    padding_values=0,
22)
23
24for batch in batched:
25    print(batch.numpy())

Using 0 as the padding value is common when 0 is reserved for padding in your vocabulary. If 0 is a real token, choose another value and keep the model’s masking logic consistent with it.

Include Labels or Multiple Fields

Training data usually includes more than just the variable-length list. You can yield tuples from the generator and declare a matching nested signature.

python
1import tensorflow as tf
2
3examples = [
4    ([1, 2, 3], 0),
5    ([4, 5], 1),
6    ([6, 7, 8, 9], 0),
7]
8
9def gen():
10    for tokens, label in examples:
11        yield tokens, label
12
13ds = tf.data.Dataset.from_generator(
14    gen,
15    output_signature=(
16        tf.TensorSpec(shape=(None,), dtype=tf.int32),
17        tf.TensorSpec(shape=(), dtype=tf.int32),
18    ),
19)
20
21batched = ds.padded_batch(
22    batch_size=2,
23    padded_shapes=((None,), ()),
24    padding_values=(0, 0),
25)
26
27for tokens, labels in batched:
28    print(tokens.numpy())
29    print(labels.numpy())

Only the sequence field is padded. The label remains a scalar for each example.

When Ragged Tensors Are a Better Fit

TensorFlow also supports ragged tensors, which preserve variable-length dimensions without padding everything up front. They are useful when downstream operations understand ragged input.

python
1import tensorflow as tf
2
3ragged = tf.ragged.constant([
4    [1, 2, 3],
5    [4, 5],
6    [6, 7, 8, 9],
7], dtype=tf.int32)
8
9ds = tf.data.Dataset.from_tensor_slices(ragged)
10
11for row in ds:
12    print(row)

Ragged tensors can be cleaner than immediate padding, especially in preprocessing pipelines. Still, not every layer or custom op supports them, so padded dense batches remain the most portable choice for many models.

Choosing Between Padding and Ragged Data

Padding is usually easier when you feed data into embedding, recurrent, or transformer-style models that already have masking support. Ragged tensors are attractive when you want to keep original sequence lengths intact for longer and avoid padding overhead during early transformations. In practice, a robust default is generator plus output signature plus padded_batch, unless you know the rest of the pipeline is ragged-aware.

Common Pitfalls

One common mistake is calling from_tensor_slices directly on a Python list of uneven inner lengths and expecting TensorFlow to infer the right structure. Another is forgetting to use shape=(None,) in the output signature, which makes the dataset too strict for variable-length rows. Developers also use batch() instead of padded_batch(), which fails because the elements have different shapes. Padding with a value that overlaps real data can quietly harm training. Finally, ragged tensors are powerful, but you should verify that every later operation in the pipeline supports them before committing to that approach.

Summary

  • Uneven nested lists cannot be represented as a normal dense tensor without extra handling.
  • 'Dataset.from_generator is a practical way to create a dataset from variable-length rows.'
  • Use tf.TensorSpec(shape=(None,), ...) to declare a one-dimensional sequence of unknown length.
  • Use padded_batch when your model needs dense batches.
  • Use ragged tensors when the rest of the pipeline can handle them cleanly.
  • Keep padding values and masking behavior aligned with the model design.

Course illustration
Course illustration

All Rights Reserved.