Introduction

In TensorFlow, you can iterate over a sequence stored in a tensor, but the best method depends on what you are trying to compute. If each step is independent, vectorized tensor operations are usually better than an explicit loop. If each step depends on earlier steps, use TensorFlow-native looping tools such as tf.unstack, tf.scan, or a tf.function loop with TensorArray.

Start with the Shape

Most sequence models store inputs in a shape like batch x time x features. Before writing any loop, identify which axis is the sequence axis.

1import tensorflow as tf
2
3x = tf.random.uniform(shape=(2, 4, 3))
4print(x.shape)

Here the time axis is 1, so iterating over the sequence means iterating over x[:, t, :] for each step t.

A surprising number of TensorFlow bugs come from looping over the batch axis by accident.

Simple Eager Iteration with tf.unstack

If you want readable code and the sequence is short, tf.unstack is the easiest starting point.

1import tensorflow as tf
2
3x = tf.random.uniform(shape=(2, 4, 3))
4steps = tf.unstack(x, axis=1)
5outputs = []
6
7for step in steps:
8    outputs.append(tf.reduce_sum(step, axis=-1))
9
10result = tf.stack(outputs, axis=1)
11print(result.shape)

This works well in eager mode and is easy to debug. The tradeoff is that it materializes many step tensors and is not ideal for long sequences or highly optimized graph execution.

Graph-Friendly Iteration with TensorArray

If you need looping inside a traced function, keep the whole operation inside TensorFlow.

1import tensorflow as tf
2
3@tf.function
4def reduce_over_time(inp):
5    time_dim = tf.shape(inp)[1]
6    out = tf.TensorArray(dtype=inp.dtype, size=time_dim)
7
8    for t in tf.range(time_dim):
9        step = inp[:, t, :]
10        out = out.write(t, tf.reduce_mean(step, axis=-1))
11
12    return tf.transpose(out.stack(), perm=[1, 0])
13
14x = tf.random.uniform(shape=(2, 4, 3))
15print(reduce_over_time(x))

This pattern is useful when the loop body cannot be expressed as one tensor operation and when you want TensorFlow to stage the computation efficiently.

Use tf.scan for True Sequential State

tf.scan is useful when each step depends on the previous state, which is common in recurrence-like logic.

1import tensorflow as tf
2
3sequence = tf.constant([1.0, 2.0, 3.0, 4.0])
4
5running_sum = tf.scan(lambda acc, x: acc + x, sequence)
6print(running_sum)

This is conceptually different from a simple map over time. tf.scan is for stateful accumulation across the sequence.

Prefer Vectorization When the Steps Are Independent

Many explicit loops are unnecessary. If the same operation is applied independently to every time step, vectorization is simpler and usually faster.

1import tensorflow as tf
2
3x = tf.random.uniform(shape=(2, 4, 3))
4result = tf.reduce_sum(x, axis=-1)
5print(result.shape)

That single operation replaces an entire loop that sums each step independently.

If you are using GPUs or TPUs, vectorized operations are usually much friendlier to the accelerator than Python-driven step iteration.

When an RNN Layer Is the Better Abstraction

If the real goal is sequence modeling rather than just manual iteration, use TensorFlow's sequence layers instead of writing your own step loop.

1import tensorflow as tf
2
3model = tf.keras.Sequential([
4    tf.keras.layers.Input(shape=(4, 3)),
5    tf.keras.layers.GRU(8),
6    tf.keras.layers.Dense(1)
7])
8
9x = tf.random.uniform(shape=(2, 4, 3))
10print(model(x).shape)

A custom loop is appropriate only when you truly need custom step logic that a built-in recurrent layer cannot express.

Common Pitfalls

The biggest mistake is using Python loops over dynamic tensors and expecting graph-level performance. That often works functionally but leaves performance on the table.

Another mistake is ignoring the sequence axis and iterating over the wrong dimension.

A third problem is writing a manual loop for a computation that is actually just one vectorized tensor reduction or one recurrent layer call.

Summary

• Use tf.unstack when you want simple eager-mode iteration over a sequence axis
• Use TensorArray and tf.function when you need a graph-friendly explicit loop
• Use tf.scan when each step depends on the previous state
• Prefer vectorized operations when each step is independent
• If the task is real sequence modeling, a built-in RNN layer is often better than manual iteration