Introduction

In TensorFlow 2, eager execution makes gradient computation much more direct than the old graph-session workflow. If you want the gradient of a model output with respect to a specific layer, the usual tool is tf.GradientTape.

The key detail is what “with respect to a layer” means. You might want gradients with respect to the layer’s output activations, or with respect to the layer’s trainable variables. Those are related but not identical questions.

Gradients with Respect to a Layer’s Activations

A common interpretation is: “how does the final output change if the activations of this intermediate layer change?” To do that, expose the intermediate activation and watch it with a gradient tape.

1import tensorflow as tf
2
3inputs = tf.keras.Input(shape=(4,))
4x = tf.keras.layers.Dense(8, activation="relu", name="hidden")(inputs)
5outputs = tf.keras.layers.Dense(1, name="out")(x)
6model = tf.keras.Model(inputs, outputs)
7
8probe_model = tf.keras.Model(
9    inputs=model.input,
10    outputs=[model.get_layer("hidden").output, model.output],
11)
12
13sample = tf.random.normal((1, 4))
14
15with tf.GradientTape() as tape:
16    hidden_output, final_output = probe_model(sample)
17    tape.watch(hidden_output)
18    target = final_output[:, 0]
19
20grads = tape.gradient(target, hidden_output)
21print(grads)

Here, hidden_output is not a variable, so tape.watch(hidden_output) is necessary. The resulting gradient tensor has the same shape as the hidden activation tensor.

Why a Probe Model Helps

The intermediate layer output is not automatically returned by a normal model call. A probe model is a convenient way to request both:

• the activation of the layer you care about
• the final output whose gradient you want

This makes the computation explicit and avoids digging through internal graph details.

Gradients with Respect to a Layer’s Weights

If instead you want the gradient of the output with respect to the layer’s trainable parameters, use the layer’s variables directly.

1import tensorflow as tf
2
3inputs = tf.keras.Input(shape=(4,))
4x = tf.keras.layers.Dense(8, activation="relu", name="hidden")(inputs)
5outputs = tf.keras.layers.Dense(1, name="out")(x)
6model = tf.keras.Model(inputs, outputs)
7
8sample = tf.random.normal((1, 4))
9hidden_layer = model.get_layer("hidden")
10
11with tf.GradientTape() as tape:
12    final_output = model(sample)
13    target = final_output[:, 0]
14
15grads = tape.gradient(target, hidden_layer.trainable_variables)
16for variable, grad in zip(hidden_layer.trainable_variables, grads):
17    print(variable.name, grad.shape)

Because trainable variables are automatically watched by the tape, you do not need tape.watch(...) in this case.

Choosing the Right Scalar Target

GradientTape.gradient expects a scalar target or something reducible to one. If the model output is a vector or matrix, choose the component you care about.

For example, in a classification model you might pick one class score:

target = logits[:, class_index]

Or reduce a tensor to a scalar:

target = tf.reduce_sum(final_output)

If you do not specify this clearly, the meaning of the gradient becomes ambiguous.

Example for a Specific Class Score

1import tensorflow as tf
2
3inputs = tf.keras.Input(shape=(6,))
4x = tf.keras.layers.Dense(10, activation="relu", name="hidden")(inputs)
5logits = tf.keras.layers.Dense(3, name="logits")(x)
6model = tf.keras.Model(inputs, logits)
7
8probe_model = tf.keras.Model(
9    inputs=model.input,
10    outputs=[model.get_layer("hidden").output, model.output],
11)
12
13sample = tf.random.normal((1, 6))
14class_index = 2
15
16with tf.GradientTape() as tape:
17    hidden_output, output_logits = probe_model(sample)
18    tape.watch(hidden_output)
19    target = output_logits[:, class_index]
20
21grad_hidden = tape.gradient(target, hidden_output)
22print(grad_hidden)

This pattern is useful for saliency-style analysis and understanding how a specific prediction depends on internal activations.

Common Pitfalls

One common mistake is asking for gradients with respect to an intermediate tensor without watching it. Non-variable tensors are not automatically tracked for differentiation unless they are watched.

Another issue is trying to differentiate a whole tensor output without deciding which scalar objective matters. Gradients are clearest when you define a specific scalar target.

It is also easy to confuse gradients with respect to activations and gradients with respect to weights. Those answer different questions and use different targets in the gradient(...) call.

Finally, if you build the intermediate tensor outside the tape context, the tape may not record the operations you need. Keep the relevant forward pass inside the GradientTape block.

Summary

• In TensorFlow 2 eager execution, tf.GradientTape is the standard tool for gradients.
• To differentiate with respect to a layer’s activations, expose the intermediate output and watch it explicitly.
• To differentiate with respect to a layer’s weights, request gradients over layer.trainable_variables.
• Choose a clear scalar target such as one output unit or a reduced summary of the output.
• Most mistakes come from not watching intermediate tensors or from mixing up activations and variables.