Practice Now
TensorFlow
2D Tensors
Element-wise Operations
Machine Learning
Deep Learning

Tensorflow apply op to each element of a 2d tensor

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
Deep LearningTensorFlow & KerasMachine Learning

Introduction

TensorFlow operations are element-wise by default for most math functions. To apply an operation to every element of a 2D tensor, use built-in ops like tf.math.square, tf.math.log, or tf.where. For custom element-wise logic, use tf.vectorized_map or tf.map_fn. Avoid Python loops over tensor elements — they bypass TensorFlow's graph execution and run orders of magnitude slower. The key principle is to express operations as vectorized tensor math wherever possible.

Built-in Element-wise Operations

python
1import tensorflow as tf
2
3t = tf.constant([[1.0, 2.0, 3.0],
4                 [4.0, 5.0, 6.0]])
5
6# Most TF math ops are already element-wise
7tf.math.square(t)      # [[1, 4, 9], [16, 25, 36]]
8tf.math.sqrt(t)        # [[1.0, 1.41, 1.73], [2.0, 2.24, 2.45]]
9tf.math.log(t)         # [[0.0, 0.69, 1.10], [1.39, 1.61, 1.79]]
10tf.math.exp(t)         # [[2.72, 7.39, 20.09], [54.60, 148.41, 403.43]]
11tf.math.abs(t)         # Absolute value
12tf.math.sin(t)         # Sine
13tf.math.sigmoid(t)     # Sigmoid activation
14
15# Arithmetic operators work element-wise
16t * 2                  # [[2, 4, 6], [8, 10, 12]]
17t + 10                 # [[11, 12, 13], [14, 15, 16]]
18t ** 2                 # Same as tf.math.square(t)
19
20# Element-wise comparison
21tf.math.greater(t, 3.0)  # [[False, False, False], [True, True, True]]

Custom Element-wise with Standard Ops

python
1# Combine built-in ops for custom element-wise functions
2t = tf.constant([[1.0, -2.0, 3.0],
3                 [-4.0, 5.0, -6.0]])
4
5# ReLU: max(0, x)
6result = tf.math.maximum(t, 0.0)  # [[1, 0, 3], [0, 5, 0]]
7
8# Clamp between -1 and 1
9result = tf.clip_by_value(t, -1.0, 1.0)  # [[-1, -1, 1], [-1, 1, -1]]
10
11# Conditional: if x > 0 then x^2 else 0
12result = tf.where(t > 0, tf.square(t), tf.zeros_like(t))
13# [[1, 0, 9], [0, 25, 0]]
14
15# Custom formula: (x - mean) / std
16mean = tf.reduce_mean(t)
17std = tf.math.reduce_std(t)
18normalized = (t - mean) / std

Using tf.map_fn for Complex Per-Element Logic

python
1# tf.map_fn applies a function to each element along the first axis
2t = tf.constant([[1.0, 2.0, 3.0],
3                 [4.0, 5.0, 6.0]])
4
5# Apply function to each row
6def process_row(row):
7    return row / tf.reduce_sum(row)  # Normalize each row
8
9result = tf.map_fn(process_row, t)
10# [[0.167, 0.333, 0.5], [0.267, 0.333, 0.4]]
11
12# For true element-wise mapping, reshape to 1D, map, reshape back
13flat = tf.reshape(t, [-1])
14
15def custom_op(x):
16    return tf.cond(x > 3.0, lambda: x * 2, lambda: x + 10)
17
18result = tf.map_fn(custom_op, flat)
19result = tf.reshape(result, t.shape)
20# [[11, 12, 13], [8, 10, 12]]

Using tf.vectorized_map (Faster)

python
1# tf.vectorized_map auto-vectorizes a function — faster than tf.map_fn
2t = tf.constant([[1.0, 2.0, 3.0],
3                 [4.0, 5.0, 6.0]])
4
5def process_row(row):
6    return tf.nn.softmax(row)
7
8result = tf.vectorized_map(process_row, t)
9# Each row is independently softmaxed
10
11# tf.vectorized_map is typically 2-10x faster than tf.map_fn
12# because it executes ops in batch rather than looping

Applying Custom Python Functions with tf.py_function

python
1import numpy as np
2
3# For truly custom logic that cannot be expressed in TF ops
4def numpy_custom_op(tensor):
5    arr = tensor.numpy()
6    # Any NumPy/Python operation
7    result = np.where(arr > 3, arr ** 0.5, arr * 2)
8    return tf.constant(result, dtype=tf.float32)
9
10t = tf.constant([[1.0, 2.0, 3.0],
11                 [4.0, 5.0, 6.0]])
12
13# tf.py_function wraps Python code for use in TF graphs
14result = tf.py_function(numpy_custom_op, [t], tf.float32)
15result.set_shape(t.shape)  # Shape inference lost — set manually
16# [[2.0, 4.0, 6.0], [2.0, 2.24, 2.45]]

Performance Comparison

python
1import time
2
3t = tf.random.normal([1000, 1000])
4
5# Fast: vectorized TF op
6start = time.time()
7r1 = tf.math.square(t) + tf.math.sin(t)
8print(f"Vectorized: {time.time() - start:.4f}s")
9
10# Moderate: tf.map_fn
11start = time.time()
12r2 = tf.map_fn(lambda row: tf.math.square(row) + tf.math.sin(row), t)
13print(f"map_fn: {time.time() - start:.4f}s")
14
15# Slow: Python loop (avoid this)
16start = time.time()
17rows = []
18for i in range(t.shape[0]):
19    rows.append(tf.math.square(t[i]) + tf.math.sin(t[i]))
20r3 = tf.stack(rows)
21print(f"Python loop: {time.time() - start:.4f}s")
22
23# Vectorized: ~0.001s, map_fn: ~0.05s, Python loop: ~2.0s

Common Pitfalls

  • Using Python loops to iterate over tensor elements: for i in range(t.shape[0]): t[i] creates separate TF operations for each element and bypasses batch optimization. Use vectorized TF ops (tf.math.square, tf.where) or tf.map_fn for operations that must apply per-row or per-element.
  • Confusing tf.map_fn axis behavior: tf.map_fn maps over the first dimension (axis 0) by default. For a 2D tensor of shape (m, n), it applies the function to each row, not each individual element. To map per-element, flatten first with tf.reshape(t, [-1]), apply, then reshape back.
  • Using tf.py_function in performance-critical code: tf.py_function drops into Python execution, losing all TF graph optimizations, GPU acceleration, and XLA compilation. It also prevents tf.function tracing. Use it only for prototyping or operations that genuinely cannot be expressed in TF ops.
  • Forgetting that TF ops broadcast automatically: t + scalar or t * vector automatically broadcasts across dimensions. Writing explicit loops or tf.map_fn for operations that are already handled by broadcasting wastes performance. Check NumPy broadcasting rules — TF follows the same conventions.
  • Shape loss with tf.py_function and tf.map_fn: Both functions may lose static shape information. TF cannot infer output shapes through Python code. Call result.set_shape(expected_shape) after tf.py_function to restore shape metadata, especially if the result feeds into a Keras layer.

Summary

  • Most TF math ops (tf.math.square, tf.where, +, *) are already element-wise — use them directly
  • Use tf.map_fn to apply a function per-row of a 2D tensor, or flatten to apply per-element
  • Prefer tf.vectorized_map over tf.map_fn for better performance (auto-batches operations)
  • Avoid Python loops over tensor elements — they are 100-1000x slower than vectorized ops
  • Use tf.py_function only as a last resort when the operation cannot be expressed in TF ops
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
Codemia
© 2026 Codemia
Resources
Legal
Social
All Rights Reserved.