GradientTape

Setup#

pip install tensorflow

Outputs verified with TensorFlow 2.21, CPU. Everything below is deterministic — no seeds needed.

From one derivative to a full training loop#

import tensorflow as tf

# 1. The smallest possible example
x = tf.Variable(3.0)
with tf.GradientTape() as tape:
    y = x ** 2
print("dy/dx at x=3:", tape.gradient(y, x).numpy())

# 2. Constants are not recorded unless you ask
c = tf.constant(3.0)
with tf.GradientTape() as tape:
    y = c ** 2
print("gradient for an unwatched constant:", tape.gradient(y, c))

with tf.GradientTape() as tape:
    tape.watch(c)
    y = c ** 2
print("after tape.watch:", tape.gradient(y, c).numpy())

# 3. A real training step, no Keras anywhere
w = tf.Variable(0.0)
b = tf.Variable(0.0)
xs = tf.constant([1.0, 2.0, 3.0, 4.0])
ys = tf.constant([3.0, 5.0, 7.0, 9.0])       # the truth is y = 2x + 1

for step in range(200):
    with tf.GradientTape() as tape:
        loss = tf.reduce_mean((w * xs + b - ys) ** 2)
    dw, db = tape.gradient(loss, [w, b])
    w.assign_sub(0.01 * dw)                   # walk downhill
    b.assign_sub(0.01 * db)

print(f"learned w={w.numpy():.2f}, b={b.numpy():.2f}, loss={loss.numpy():.4f}")

dy/dx at x=3: 6.0
gradient for an unwatched constant: None
after tape.watch: 6.0
learned w=2.05, b=0.84, loss=0.0042

The walkthrough#

6.0 is the calculus checking out: the derivative of x² is 2x, and 2 × 3 = 6. When learning autodiff, always start with a function whose derivative you can verify on paper.

tf.Variable vs tf.constant is the recording rule. Variables are watched automatically — they are the "knobs" training exists to turn. Constants are treated as fixed scenery and ignored, which is why the unwatched gradient came back None rather than zero. None means never recorded; zero would mean recorded, no influence. The distinction diagnoses many bugs.

The training loop is fit with the covers off. Loss inside the tape, tape.gradient(loss, [w, b]) for the blame, assign_sub to walk downhill with step size 0.01. Two hundred steps land near the truth (w=2, b=1); the remaining gap closes with more steps or a schedule — the walk itself is gradient descent.

One tape, one replay. After tape.gradient returns, the tape frees its recording. Ask twice and you get: RuntimeError: A non-persistent GradientTape can only be used to compute one set of gradients. When you genuinely need two replays — two different outputs, one recording — create it as tf.GradientTape(persistent=True) and del tape when done.

Common mistakes#

Computing outside the with block. Any operation above the with line is off the record, and its gradient path is dead. The forward pass — all of it — goes inside.

Breaking the chain with .numpy(). Converting a tensor to a NumPy value mid-calculation leaves the tape's world; everything downstream is invisible to it. Stay in tensors until the loss is computed. The same rule exists in PyTorch as detach.

Expecting gradients through integers. Gradients flow through float tensors only. An int32 input silently yields None — check dtypes when a gradient is mysteriously missing.

Reassigning instead of assign_sub. Writing w = w - 0.01 * dw rebinds the Python name to a new tensor, and the original Variable never changes. In-place updates on Variables go through assign, assign_add, assign_sub.

Try it yourself#

Replace the loss with mean absolute error: tf.reduce_mean(tf.abs(w * xs + b - ys)). Rerun and watch the loop still work — the tape differentiates the new expression without any other change. That indifference to what you wrote is the whole point of automatic differentiation.

What to learn next#

• tf.function, graphs and retracing — making tape-based loops fast.
• Backpropagation — the algorithm the tape implements.
• Autograd graph mechanics — the identical machinery in PyTorch.