tensorflow
masking
sequence length
deep learning
data preprocessing

tensorflow creating mask of varied lengths

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

Masks are how TensorFlow tells a model which time steps are real data and which time steps are only padding. When sequences in a batch have different lengths, creating the correct mask is essential because the model should ignore padded positions during attention, recurrent processing, or loss computation.

Create a Mask From Sequence Lengths

If you already know the valid length of each sequence in the batch, tf.sequence_mask is the most direct solution.

python
1import tensorflow as tf
2
3lengths = tf.constant([3, 1, 4])
4mask = tf.sequence_mask(lengths, maxlen=5)
5print(mask.numpy())

This produces a boolean mask where True means "real token" and False means "padding." The first sequence marks three valid positions, the second marks one, and the third marks four.

This is the standard answer when your preprocessing pipeline already tracks sequence lengths separately.

Create a Mask From Padded Data

If the data is padded with a known value such as 0, you can derive the mask directly from the tensor instead of storing lengths in a separate array.

python
1import tensorflow as tf
2
3sequences = tf.constant([
4    [7, 2, 5, 0, 0],
5    [3, 0, 0, 0, 0],
6    [9, 8, 1, 4, 0],
7])
8
9mask = tf.not_equal(sequences, 0)
10print(mask.numpy())

This is especially common in NLP pipelines where token ID 0 is reserved for padding.

The choice between tf.sequence_mask and tf.not_equal depends on where the truth already lives. If lengths are known, use lengths. If padded values are reliable, deriving the mask from the tensor can be simpler.

Use Keras Masking Support When Possible

Many Keras layers can propagate masks automatically. For padded sequence inputs, Embedding(mask_zero=True) is often the cleanest solution.

python
1import tensorflow as tf
2
3model = tf.keras.Sequential([
4    tf.keras.layers.Embedding(input_dim=20, output_dim=8, mask_zero=True),
5    tf.keras.layers.LSTM(16),
6    tf.keras.layers.Dense(1)
7])
8
9x = tf.constant([
10    [1, 2, 3, 0, 0],
11    [4, 5, 0, 0, 0],
12    [6, 7, 8, 9, 0],
13])
14
15print(model(x))

Here, the embedding layer creates the mask automatically and passes it to the LSTM. That is usually better than manually threading masks through every layer when the built-in masking system already fits the architecture.

Match Mask Shape to the Operation

Different TensorFlow operations expect masks in different shapes. Sequence models often use a two-dimensional mask of shape batch x time, while attention mechanisms may need the mask expanded to additional dimensions.

For example, turning a sequence mask into a float attention mask can look like this:

python
1import tensorflow as tf
2
3lengths = tf.constant([3, 2])
4mask = tf.sequence_mask(lengths, maxlen=4)
5attention_mask = tf.cast(mask[:, tf.newaxis, tf.newaxis, :], tf.float32)
6print(attention_mask.shape)

The underlying information is the same, but the shape is adapted for the operation consuming it.

Use the Mask in Losses or Metrics When Needed

Sometimes the layer stack handles masking automatically, but the loss function still needs manual masking. That happens in token-level tasks where padded positions should not contribute to the objective.

python
1import tensorflow as tf
2
3losses = tf.constant([[0.3, 0.2, 0.1, 0.0], [0.5, 0.4, 0.0, 0.0]])
4mask = tf.constant([[1, 1, 1, 0], [1, 1, 0, 0]], dtype=tf.float32)
5
6masked_loss = tf.reduce_sum(losses * mask) / tf.reduce_sum(mask)
7print(masked_loss.numpy())

That pattern keeps padded positions from distorting the reported training signal.

Common Pitfalls

A common mistake is creating the right mask values but the wrong dtype or shape for the downstream operation. Always check what the target layer expects.

Another mistake is assuming every Keras layer will automatically propagate masks. Many sequence layers do, but custom layers and some lower-level TensorFlow operations will not unless you handle masking yourself.

People also often forget that the padding token must be reserved consistently. If 0 sometimes means padding and sometimes means a real token, a derived mask from tf.not_equal(..., 0) becomes invalid.

Finally, do not let padded positions leak into losses or metrics when the model architecture does not mask them automatically.

Summary

  • Use tf.sequence_mask when you have explicit sequence lengths.
  • Use comparisons such as tf.not_equal(x, 0) when padded values are known and reliable.
  • Let Keras propagate masks automatically when layers such as Embedding(mask_zero=True) fit the model.
  • Adapt the mask shape to the operation that consumes it.
  • Remember that masking sometimes needs to be applied in the loss, not just in the model layers.

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.