TensorFlow
Ragged Tensors
Batch Processing
Machine Learning
TensorFlow 2.0

How do I make a ragged batch in Tensorflow 2.0?

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design

Introduction

A ragged batch in TensorFlow is useful when each sample has variable length, such as token sequences, variable-size detection boxes, or uneven event logs. If you force padding too early, you may waste memory and distort sequence statistics. tf.RaggedTensor keeps variable-length dimensions explicit while still enabling batching and model operations that support ragged inputs.

In TensorFlow 2.x, you can build ragged batches through tf.ragged.constant or tf.data.Dataset.ragged_batch depending on pipeline style.

Core Sections

1. Create ragged tensors directly

python
1import tensorflow as tf
2
3samples = [
4    [1, 2, 3],
5    [4, 5],
6    [6, 7, 8, 9]
7]
8
9rt = tf.ragged.constant(samples)
10print(rt)

Shape becomes [3, None] where second dimension is variable.

2. Build ragged batches with tf.data

python
1def gen():
2    yield [1, 2, 3]
3    yield [4, 5]
4    yield [6, 7, 8, 9]
5
6ds = tf.data.Dataset.from_generator(
7    gen,
8    output_signature=tf.TensorSpec(shape=(None,), dtype=tf.int32)
9)
10
11ragged_ds = ds.ragged_batch(2)
12for batch in ragged_ds:
13    print(batch)

ragged_batch preserves variable lengths within batches.

3. Convert ragged to dense when required

Some layers need dense tensors. Convert with padding explicitly:

python
dense = rt.to_tensor(default_value=0)
print(dense)

Choose padding value and mask handling carefully.

4. Use compatible layers

Many TensorFlow ops support ragged inputs, but not all. For sequence models, embeddings and some recurrent layers can work with ragged or masked dense alternatives. Verify layer docs before committing architecture.

5. Track row lengths for downstream logic

python
lengths = rt.row_lengths()
print(lengths)

Useful for custom losses, masking, or attention constraints.

Common Pitfalls

  • Padding everything upfront and losing ragged efficiency advantages.
  • Using standard batch instead of ragged_batch for variable-length samples.
  • Passing ragged tensors into layers that only accept dense inputs.
  • Converting to dense without explicit masking strategy.
  • Ignoring row lengths and misaligning sequence-level computations.

Summary

To make ragged batches in TensorFlow 2.x, represent variable-length data with tf.RaggedTensor and batch using ragged_batch in tf.data pipelines. Convert to dense only where necessary, with explicit padding and masking logic. With this approach, variable-length workloads remain efficient and semantically correct throughout training and inference.

A practical way to make this guidance durable is to convert it into a small runbook that includes prerequisites, expected environment versions, and a short verification sequence. Even strong teams lose time when troubleshooting steps live only in memory or chat history. A runbook should explicitly answer three questions: what to check first, what output confirms healthy behavior, and what output indicates a known failure mode. This level of clarity helps both experienced maintainers and newer contributors, and it reduces repeated investigation during incidents.

It is also valuable to create a tiny reproducible fixture for this topic. The fixture can be a minimal script, test case, sample request, or small dataset that demonstrates the correct behavior in isolation. When regressions appear after dependency upgrades, infrastructure changes, or framework migrations, that fixture becomes the fastest way to isolate whether the issue is environmental or logic-related. Keeping a focused fixture in source control gives you a stable benchmark across branches and release cycles.

For long-term reliability, pair documentation with one automated guardrail in CI. The guardrail should be narrow and fast: an import check, schema validation, endpoint contract test, deterministic unit test, or lightweight performance threshold. Avoid broad flaky checks that hide real signals. The goal is early, actionable feedback before code reaches production. If the same category of issue appears repeatedly, promote the manual troubleshooting step into automation so the system catches it first. Over time, this shifts effort from reactive debugging to preventive quality control and keeps the knowledge article relevant in real engineering workflows.

As a final hardening step, periodically rerun the sample code in a clean environment image and record results in version control. This catches ecosystem drift early and keeps implementation guidance aligned with real runtime behavior.


Related reading
Free course
Beginner
7 lessons
2 hours
Tackling System Design Interview Problems

A short course that equips you with the skills to approach system design interviews methodically.

Start the free course
Track what you have practised

A free account saves your progress, solutions and study plan across every problem on Codemia.

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design

All Rights Reserved.