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

FunctionDescription
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)
tf.math.exp(x)
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, tanTrig
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.print for 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])