Introduction

You can loop over tensors in libraries such as PyTorch and TensorFlow, but whether you should is a different question. In many numeric workloads, explicit Python loops are much slower than vectorized tensor operations. The right approach depends on whether you need readability, per-element control, or maximum throughput.

What Looping Over a Tensor Means

A tensor is just a multi-dimensional array, so iterating over it means stepping through one dimension at a time. In Python, iteration over a rank-2 tensor usually yields rows. Iteration over a rank-1 tensor yields scalar elements or small tensor objects depending on the library.

PyTorch example:

1import torch
2
3x = torch.tensor([[1, 2], [3, 4], [5, 6]])
4
5for row in x:
6    print(row)

This prints one row tensor at a time. That is perfectly legal, but it is still Python-side iteration.

PyTorch: Legal but Often Slow

PyTorch allows straightforward loops, which is helpful for debugging or custom logic.

1import torch
2
3x = torch.tensor([1.0, 2.0, 3.0, 4.0])
4total = 0.0
5
6for value in x:
7    total += float(value)
8
9print(total)

That works, but the vectorized form is better:

1import torch
2
3x = torch.tensor([1.0, 2.0, 3.0, 4.0])
4total = torch.sum(x)
5print(total.item())

The second version keeps the work in optimized tensor kernels instead of bouncing between Python and the tensor runtime.

Use explicit loops in PyTorch when:

• you are debugging shapes or intermediate values
• the loop count is tiny and clarity matters more than speed
• each step truly depends on Python-side control flow

Otherwise, prefer tensor ops.

TensorFlow: Eager Mode Versus Graph Mode

TensorFlow makes this distinction even sharper. In eager mode, you can iterate over tensors more naturally. In graph-oriented code, plain Python loops may not behave the way you expect because the graph needs symbolic operations.

Eager example:

1import tensorflow as tf
2
3x = tf.constant([[1, 2], [3, 4], [5, 6]])
4
5for row in x:
6    print(row.numpy())

If you need a TensorFlow-native loop in graph-style code, use tf.map_fn or tf.while_loop instead of depending on Python iteration.

1import tensorflow as tf
2
3x = tf.constant([1, 2, 3, 4])
4y = tf.map_fn(lambda v: v * v, x)
5print(y.numpy())

This matters when building models or functions that will be traced with tf.function.

Prefer Vectorization When Possible

Most tensor looping questions are really optimization questions. A loop that applies the same formula element by element should usually be written as one tensor expression.

Slow style:

1import torch
2
3x = torch.arange(5, dtype=torch.float32)
4result = []
5
6for value in x:
7    result.append(value * value + 1)
8
9print(torch.stack(result))

Better style:

1import torch
2
3x = torch.arange(5, dtype=torch.float32)
4result = x * x + 1
5print(result)

The vectorized version is shorter, easier to optimize, and more GPU-friendly.

Loop Over Batches, Not Elements

If you do need loops, loop at a higher level. Iterating over batches or sequences is usually much more reasonable than iterating over scalar tensor elements.

1import torch
2
3batches = torch.tensor(
4    [
5        [1.0, 2.0, 3.0],
6        [4.0, 5.0, 6.0],
7    ]
8)
9
10for batch in batches:
11    print(batch.mean())

This pattern shows up in custom training or sequence processing, and it keeps the expensive work inside tensor operations even though the outer control flow uses Python.

When Explicit Tensor Loops Are Justified

Explicit looping is still reasonable when:

• the algorithm is inherently sequential
• each step depends on the previous step's result
• you are implementing a prototype before optimizing
• the tensor is small and performance is irrelevant

Examples include beam-search style decoding, certain dynamic programs, or debugging code where correctness comes first.

Common Pitfalls

• Iterating element by element when a vectorized operation exists.
• Using Python loops in TensorFlow graph-style code where tf.map_fn or tf.while_loop is required.
• Moving tensors back to Python scalars inside hot loops and destroying performance.
• Assuming iteration over a tensor yields scalars when it may yield sub-tensors.
• Writing GPU code that spends most of its time in Python control flow instead of tensor kernels.

Summary

• You can loop over tensors, but it is often not the best approach.
• PyTorch allows direct iteration more naturally, while TensorFlow often needs graph-aware looping primitives.
• Prefer vectorized tensor expressions whenever the operation is uniform.
• If you must loop, loop over larger units such as rows, batches, or time steps.
• Treat explicit per-element loops as a correctness or debugging tool, not the default numeric strategy.