Gradients for Many Variables with tf.GradientTape

The Story — The Sound Engineer's Mixing Desk

In the last lesson you found the gradient for one variable. Real models have many. The good news: the tape handles any number of variables at once, with one backward pass.

Two Variables: A Line y = w·x + b

We want the line to predict y = 7 when x = 2. Our current guess is w = 1, b = 0. The loss is the squared error.

import tensorflow as tf

w = tf.Variable(1.0)
b = tf.Variable(0.0)
x, target = 2.0, 7.0

with tf.GradientTape() as tape:
    pred = w * x + b
    loss = (pred - target) ** 2

dw, db = tape.gradient(loss, [w, b])      # a list in, a list out
print("pred :", pred.numpy(), " loss:", loss.numpy())
print("dL/dw:", dw.numpy())
print("dL/db:", db.numpy())

Any Structure Works: List, Tuple or Dict

The gradient comes back in the same structure you pass in.

with tf.GradientTape() as tape:
    loss = (w * x + b - target) ** 2

grads = tape.gradient(loss, {"weight": w, "bias": b})
print({k: float(v) for k, v in grads.items()})

Try It — Walk on the Loss Landscape

Gradients Have the Same Shape as Variables

This is a key rule. A weight matrix of shape (3, 2) gets a gradient of shape (3, 2). One number per knob.

tf.random.set_seed(1)
X = tf.random.normal([5, 3])                 # batch of 5, 3 features
y = tf.random.normal([5, 2])                 # 2 targets
W = tf.Variable(tf.random.normal([3, 2]))
b = tf.Variable(tf.zeros([2]))

with tf.GradientTape() as tape:
    pred = X @ W + b
    loss = tf.reduce_mean((pred - y) ** 2)

dW, db = tape.gradient(loss, [W, b])
print("W", W.shape, "-> dW", dW.shape)
print("b", b.shape, "   -> db", db.shape)

A Whole Keras Model in One Call

model = tf.keras.Sequential([
    tf.keras.Input(shape=(3,)),
    tf.keras.layers.Dense(4, activation="relu"),
    tf.keras.layers.Dense(1),
])

with tf.GradientTape() as tape:
    loss = tf.reduce_mean(model(X) ** 2)

grads = tape.gradient(loss, model.trainable_variables)
for var, g in zip(model.trainable_variables, grads):
    print(f"{var.path:26s} {str(var.shape):8s} grad {g.shape}")

Try It — Dense Layer Shape Builder

Two Surprises to Know

1. A Non-scalar Target Is Summed First

If the target is a vector, tape.gradient gives the gradient of its sum. For one gradient per output, use tape.jacobian.

x = tf.Variable(2.0)
with tf.GradientTape(persistent=True) as tape:
    ys = tf.stack([x, x ** 2, x ** 3])      # [2, 4, 8]

print("gradient :", tape.gradient(ys, x).numpy(), "  (1 + 4 + 12)")
print("jacobian :", tape.jacobian(ys, x).numpy())
del tape

2. An Unused Variable Gives None

a = tf.Variable(1.0)
unused = tf.Variable(5.0)

with tf.GradientTape() as tape:
    loss = a * 3

print(tape.gradient(loss, [a, unused]))

with tf.GradientTape() as tape:
    loss = a * 3
print(tape.gradient(loss, [a, unused],
      unconnected_gradients=tf.UnconnectedGradients.ZERO))

Golden Rules