TensorFlow Cheatsheet
Operations
Use this TensorFlow reference while you build software engineering projects, review code for technical interview prep, or polish examples for a software engineer resume.
Element-wise Math
All standard Python operators (+, -, *, /, **, //, %) are overloaded on tensors and call the corresponding tf.math.* function.
import tensorflow as tf a = tf.constant([1.0, 2.0, 3.0]) b = tf.constant([4.0, 5.0, 6.0]) a + b # tf.math.add a - b # tf.math.subtract a * b # tf.math.multiply a / b # tf.math.divide (true division) a ** 2 # tf.math.pow a // b # tf.math.floordiv a % b # tf.math.floormod
Common tf.math Functions
| Function | Description |
|---|---|
tf.math.abs(x) | Absolute value |
tf.math.negative(x) | Negate |
tf.math.sign(x) | -1, 0, or 1 |
tf.math.sqrt(x) | Square root (float) |
tf.math.square(x) | x² |
tf.math.exp(x) | eˣ |
tf.math.log(x) | Natural log |
tf.math.log1p(x) | log(1+x), numerically stable |
tf.math.sigmoid(x) | 1/(1+e⁻ˣ) |
tf.math.softplus(x) | log(1+eˣ) |
tf.math.ceil(x) | Ceiling |
tf.math.floor(x) | Floor |
tf.math.round(x) | Round to nearest even |
tf.math.minimum(x, y) | Element-wise min |
tf.math.maximum(x, y) | Element-wise max |
tf.math.clip_by_value(x, min, max) | Clip |
tf.math.sin(x), cos, tan | Trig |
tf.math.atan2(y, x) | arctan2 |
tf.math.erf(x) | Error function |
tf.math.sqrt(tf.constant([4.0, 9.0, 16.0])) # [2., 3., 4.] tf.clip_by_value(a, 1.5, 2.5) # [1.5, 2., 2.5]
Reduction Operations
t = tf.constant([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]) tf.reduce_sum(t) # 21.0 (all elements) tf.reduce_sum(t, axis=0) # [5., 7., 9.] (collapse rows) tf.reduce_sum(t, axis=1) # [6., 15.] (collapse cols) tf.reduce_sum(t, axis=1, keepdims=True) # [[6.], [15.]] tf.reduce_mean(t) tf.reduce_min(t) tf.reduce_max(t) tf.reduce_prod(t) # Logical reductions (bool tensors) b = tf.constant([[True, False], [True, True]]) tf.reduce_all(b) # False tf.reduce_any(b) # True tf.reduce_all(b, axis=0) # [True, False]
Cumulative
tf.cumsum(tf.constant([1, 2, 3, 4])) # [1, 3, 6, 10] tf.cumsum(t, axis=1) tf.cumprod(tf.constant([1, 2, 3, 4])) # [1, 2, 6, 24]
Linear Algebra
a = tf.constant([[1.0, 2.0], [3.0, 4.0]]) b = tf.constant([[5.0, 6.0], [7.0, 8.0]]) tf.linalg.matmul(a, b) # or a @ b tf.matmul(a, b) # same tf.matmul(a, b, transpose_b=True) tf.matmul(a, b, adjoint_b=True) # conjugate transpose tf.linalg.matvec(a, tf.constant([1.0, 2.0])) # matrix-vector tf.tensordot(a, b, axes=1) tf.linalg.det(a) tf.linalg.inv(a) tf.linalg.trace(a) tf.linalg.diag(tf.constant([1.0, 2.0, 3.0])) tf.linalg.band_part(a, 0, -1) # upper triangle
Decompositions
# Eigendecomposition (symmetric) vals, vecs = tf.linalg.eigh(a) # SVD s, u, v = tf.linalg.svd(a) # Cholesky (positive definite) tf.linalg.cholesky(a) # QR q, r = tf.linalg.qr(a) # Solve Ax = b x = tf.linalg.solve(a, b)
Broadcasting
TensorFlow follows NumPy broadcasting rules: dimensions are compared right-to-left; a dimension of 1 stretches to match the other.
a = tf.constant([[1.0], [2.0], [3.0]]) # (3, 1) b = tf.constant([10.0, 20.0]) # (2,) a + b # (3, 2) — broadcast # Check compatibility tf.broadcast_to(b, [3, 2]) tf.broadcast_dynamic_shape([3, 1], [2]) # → [3, 2]
Concatenation and Stacking
a = tf.constant([[1, 2], [3, 4]]) b = tf.constant([[5, 6], [7, 8]]) tf.concat([a, b], axis=0) # stack rows → (4, 2) tf.concat([a, b], axis=1) # stack cols → (2, 4) tf.stack([a, b]) # new axis 0 → (2, 2, 2) tf.stack([a, b], axis=1) # new axis 1 → (2, 2, 2) tf.unstack(a, axis=0) # list of rows
Splitting
t = tf.constant([[1, 2, 3, 4]]) tf.split(t, 2, axis=1) # two (1,2) tensors tf.split(t, [1, 3], axis=1) # sizes 1 and 3 # Alias tf.unstack(tf.constant([1, 2, 3])) # [1, 2, 3] as scalars
Sorting and Searching
t = tf.constant([3, 1, 4, 1, 5, 9, 2, 6]) tf.sort(t) # ascending tf.sort(t, direction='DESCENDING') tf.argsort(t) # indices that sort tf.argsort(t, direction='DESCENDING') tf.math.top_k(t, k=3) # values & indices of top 3 # Argmax / argmin m = tf.constant([[3, 1], [4, 2]]) tf.argmax(m, axis=0) # [1, 1] tf.argmin(m, axis=1) # [1, 1] # Where tf.where(tf.constant([True, False, True]), [1, 2, 3], [10, 20, 30]) # → [1, 20, 3] tf.where(t > 3) # indices of matching elements
Comparison Operations
a = tf.constant([1, 2, 3]) b = tf.constant([2, 2, 2]) tf.equal(a, b) # [False, True, False] tf.not_equal(a, b) tf.less(a, b) tf.less_equal(a, b) tf.greater(a, b) tf.greater_equal(a, b) # Logical tf.logical_and(tf.constant([True, False]), tf.constant([True, True])) tf.logical_or(...) tf.logical_not(...) tf.logical_xor(...)
tf.function — Graph Compilation
Decorating a Python function with @tf.function traces it into a TensorFlow graph for performance.
@tf.function def step(x): return tf.matmul(x, x) step(tf.ones([100, 100])) # first call traces; subsequent calls use graph # Specify input signatures to avoid retracing @tf.function(input_signature=[tf.TensorSpec([None, 2], tf.float32)]) def process(x): return tf.reduce_sum(x, axis=1)
Gotcha: Python side-effects (print, list appends) only run during tracing, not on every call. Use
tf.printfor in-graph printing.
@tf.function def debug_fn(x): tf.print("x =", x) # runs every call print("tracing") # runs only during trace return x + 1
Gradient Tape
x = tf.Variable(3.0) with tf.GradientTape() as tape: y = x ** 2 + 2 * x + 1 dy_dx = tape.gradient(y, x) # 2x + 2 = 8.0 # Multiple gradients with tf.GradientTape() as tape: y = x ** 3 grads = tape.gradient(y, [x]) # Higher-order gradients with tf.GradientTape() as outer: with tf.GradientTape() as inner: y = x ** 2 dy = inner.gradient(y, x) d2y = outer.gradient(dy, x) # Watch non-Variable tensors c = tf.constant(2.0) with tf.GradientTape() as tape: tape.watch(c) y = c ** 2 tape.gradient(y, c) # 4.0 # Persistent tape (multiple gradient calls) with tf.GradientTape(persistent=True) as tape: y = x ** 2 z = x ** 3 tape.gradient(y, x) tape.gradient(z, x) del tape
Padding and Tiling
t = tf.constant([[1, 2], [3, 4]]) tf.pad(t, [[1, 1], [2, 2]]) # pad rows by 1 top/bottom, cols by 2 each tf.pad(t, [[0, 1], [0, 0]], constant_values=-1) tf.tile(t, [2, 3]) # repeat 2x row-wise, 3x col-wise
Scatter and Update Ops
# Scatter (build sparse → dense) indices = tf.constant([[0], [2]]) updates = tf.constant([10.0, 30.0]) shape = tf.constant([4]) tf.scatter_nd(indices, updates, shape) # [10., 0., 30., 0.] # Update a variable at indices var = tf.Variable([0.0, 0.0, 0.0, 0.0]) var.scatter_nd_update([[1], [3]], [5.0, 7.0])