How to calculate gradients using tf.GradientTape in TensorFlow in Python

TensorFlow’s tf.GradientTape is a powerful tool for automatic differentiation, allowing you to record operations for computing gradients later. The essence of it’s straightforward: you run your forward pass inside the context of a gradient tape, and then TensorFlow tracks every computation involving variables. Once you exit that context, you can call tape.gradient() to get the derivatives with respect to any inputs.

Under the hood, the tape records a graph of operations executed during the forward pass. This graph isn’t the same static graph you might think of from TensorFlow 1.x; instead, it’s a dynamic graph built on-the-fly, capturing precisely what you’ve computed. When you request gradients, TensorFlow performs reverse-mode differentiation—walking this graph backwards to calculate gradients efficiently.

One subtlety to keep in mind is that GradientTape by default is “persistent” set to False, which means the tape is erased once gradient() is invoked. If you need to compute multiple gradients over the same forward pass, you have to instantiate the tape with persistent=True. This comes with a memory overhead, since the internal state of the tape must be preserved for each gradient operation.

Here’s a minimal example that captures the essence:

import tensorflow as tf

x = tf.Variable(3.0)  # Initialize variable

with tf.GradientTape() as tape:
    y = x * x * x  # Compute x³

# Compute dy/dx
dy_dx = tape.gradient(y, x)
print(dy_dx.numpy())  # Outputs 27.0, which is 3 * 3²

Notice how the computation y = x * x * x is executed inside the tape context. That is the critical part: if you move computation outside the context, TensorFlow won’t track it. Also, the x must be a tf.Variable or have the tf.Variable wrapper-like behavior, since gradients are defined with respect to variables and not just raw tensors.

There’s also the question of watching tensors manually. Sometimes you want to compute gradients with respect to things that aren’t variables. You can do this by explicitly telling the tape to watch a tensor using tape.watch(tensor). For example:

a = tf.constant(2.0)
with tf.GradientTape() as tape:
    tape.watch(a)  # Start watching the constant tensor
    b = a * a + 3

db_da = tape.gradient(b, a)  # Computes gradient with respect to 'a'
print(db_da.numpy())  # Should print 4.0 (2 * a)

Behind the scenes, that is important because tf.constant does not get watched automatically. Only variables are watched by default since those are supposed to be the parameters you optimize.

Another point is nested tapes. When you want to compute higher-order derivatives, you can nest tapes, one inside the other. Consider you want the second derivative:

x = tf.Variable(5.0)

with tf.GradientTape() as outer_tape:
    with tf.GradientTape() as inner_tape:
        y = x * x
    dy_dx = inner_tape.gradient(y, x)  # First derivative, 2x
d2y_dx2 = outer_tape.gradient(dy_dx, x)  # Second derivative, 2

print(dy_dx.numpy())    # 10.0
print(d2y_dx2.numpy())  # 2.0

This double tape arrangement tells TensorFlow to track all operations during both forward inner tape and the outer tape, enabling automated computation of higher derivatives without manual algebraic intervention.

Underneath, TensorFlow dynamically generates derivative expressions of each primitive operation (add, multiply, sin, cos, etc.) at runtime. This means you’re not limited to standard functions; any differentiable function comprised of TensorFlow ops will yield gradients automatically, even complex combinations involving control flow or conditional operations within the tape context.

If memory efficiency is critical, consider the difference between persistent=True and default behavior again. Persistent tapes keep the graph around for multiple gradient calls, but you’ll need to delete them manually to avoid leaks:

x = tf.Variable(4.0)

with tf.GradientTape(persistent=True) as tape:
    y = x * x * x

dy1_dx = tape.gradient(y, x)
dy2_dx = tape.gradient(y, x)  # Same tape reused

del tape  # Important to release resources

Summarizing the mechanics: place forward computations in the tape context, access gradients post-context, manage resources especially when reusing tapes, and watch any non-variable tensors explicitly. This pattern underpins nearly all custom training loops when you need gradients for optimization or analytical insight in TensorFlow.

Even more, the taping mechanism is flexible enough for mixed forward/backward passes, custom layers, and conditional control flow, ensuring that the gradient calculation precisely matches the executed computation path. It’s this dynamism that makes tf.GradientTape so much more than a convenience; it’s the toolkit that elevates TensorFlow beyond static graph constraints and into the adaptive domain-machine learning demands.

Continuing from understanding these internal mechanics, the next step is applying tf.GradientTape in real-world scenarios where gradients aren’t just theoretical but form the backbone of model training loops, optimization strategies, and parameter tuning routines. Let’s look at how practical implementations take advantage of this system without sacrificing clarity or efficiency—

Implementing gradient calculations in practical scenarios

Consider a simpler linear regression example where you want to fit a model y = Wx + b to some data points. Instead of relying on pre-built training loops, you implement the gradient calculation manually to have full control over every step:

import tensorflow as tf

# Sample data: y = 3x + 2 with some noise
x_data = tf.constant([1.0, 2.0, 3.0, 4.0])
y_data = tf.constant([5.0, 8.0, 11.0, 14.0])

# Variables to optimize
W = tf.Variable(0.0)
b = tf.Variable(0.0)

learning_rate = 0.01
epochs = 1000

for epoch in range(epochs):
    with tf.GradientTape() as tape:
        y_pred = W * x_data + b  # Forward pass
        loss = tf.reduce_mean(tf.square(y_data - y_pred))  # MSE loss

    # Compute gradients of loss wrt W and b
    gradients = tape.gradient(loss, [W, b])
    # Gradient descent step
    W.assign_sub(learning_rate * gradients[0])
    b.assign_sub(learning_rate * gradients[1])

    if epoch % 100 == 0:
        print(f"Epoch {epoch}: Loss = {loss.numpy():.4f}, W = {W.numpy():.4f}, b = {b.numpy():.4f}")

Here, the gradient tape records the computation of the loss for the given W and b. Calling tape.gradient(loss, [W, b]) instantly gives the vector of gradients corresponding to each variable, neatly packaging what would otherwise require manual symbolic derivatives.

This example also demonstrates how gradient calculation intertwines with model updating. TensorFlow variables update in-place with assign_sub(), effectively performing a manual optimization step per iteration.

Extending this to more complex operations, suppose you want to train a simple neural network layer manually without Keras conveniences. You can create a dense layer and compute gradients for weights and bias like this:

import tensorflow as tf

# Inputs and targets
x = tf.constant([[1.0, 2.0], [3.0, 4.0]])
y_true = tf.constant([[1.0], [0.0]])

# Initialize weights and bias
weights = tf.Variable(tf.random.normal([2, 1]))
bias = tf.Variable(tf.zeros([1]))

learning_rate = 0.1

for epoch in range(200):
    with tf.GradientTape() as tape:
        logits = tf.matmul(x, weights) + bias
        loss = tf.reduce_mean(tf.nn.sigmoid_cross_entropy_with_logits(labels=y_true, logits=logits))

    grads = tape.gradient(loss, [weights, bias])
    weights.assign_sub(learning_rate * grads[0])
    bias.assign_sub(learning_rate * grads[1])

    if epoch % 50 == 0:
        print(f"Epoch {epoch}: Loss = {loss.numpy():.4f}")

This example highlights that gradients propagate seamlessly through matrix operations, activations (like sigmoid cross-entropy here), and bias addition, all within the tape context. You don’t need to write derivatives by hand, regardless of the computational complexity.

When integrating with custom loss functions or applying conditions, encapsulating all related computations within the gradient tape ensures gradients reflect the exact path taken. For example, consider a hinge loss function customized for binary classification:

def hinge_loss(y_true, y_pred):
    return tf.reduce_mean(tf.maximum(0.0, 1 - y_true * y_pred))

x = tf.constant([[2.0], [-1.5], [0.0]])
y_true = tf.constant([1.0, -1.0, 1.0])

W = tf.Variable([[1.0]])
b = tf.Variable(0.0)
learning_rate = 0.01

for _ in range(300):
    with tf.GradientTape() as tape:
        logits = tf.matmul(x, W) + b
        loss = hinge_loss(y_true, tf.squeeze(logits))

    grads = tape.gradient(loss, [W, b])
    W.assign_sub(learning_rate * grads[0])
    b.assign_sub(learning_rate * grads[1])

Notice how even the control flow embedded inside the maximum operation is differentiable because it’s part of TensorFlow’s differentiable ops. The tape efficiently computes subgradients where needed.

In reinforcement learning or custom optimization algorithms, gradient tapes can also be used for policy gradient methods. For instance, given a scalar objective J(θ), gradients can be computed with respect to policy parameters θ directly:

theta = tf.Variable([0.5, -0.5])  # Policy parameters

def policy_action_prob(params):
    return tf.nn.softmax(params)

def objective(params):
    probs = policy_action_prob(params)
    reward = tf.constant([1.0, 2.0])
    return tf.reduce_sum(probs * reward)  # Expected reward

with tf.GradientTape() as tape:
    loss = -objective(theta)  # Maximize objective → minimize -objective

grad = tape.gradient(loss, theta)
theta.assign_sub(0.1 * grad)

print(theta.numpy())

The tape captures the entire flow from parameters through softmax to expected reward, enabling gradient-based optimization even in probabilistic or expectation-driven functions.

Lastly, when dealing with complex models or recursive operations where performance matters, you can combine tf.function decorators with gradient tapes to gain both speed and flexibility:

@tf.function
def train_step(x, y, W, b):
    with tf.GradientTape() as tape:
        y_pred = tf.matmul(x, W) + b
        loss = tf.reduce_mean(tf.square(y - y_pred))
    grads = tape.gradient(loss, [W, b])
    return loss, grads

W = tf.Variable(tf.random.normal([3, 1]))
b = tf.Variable(tf.zeros([1]))

x = tf.random.normal([10, 3])
y = tf.random.normal([10, 1])

for _ in range(100):
    loss, grads = train_step(x, y, W, b)
    W.assign_sub(0.01 * grads[0])
    b.assign_sub(0.01 * grads[1])

print(f"Loss after training: {loss.numpy():.4f}")

Employing tf.function here compiles the function into a static graph after the first run, greatly improving runtime without losing the automatic differentiation benefits. Still, all gradient logic remains explicit and transparent.