Introduction

In TensorFlow 2.x with Keras, gradients are typically computed using tf.GradientTape. Unlike TF1 graph-style APIs, TF2 relies on eager execution and tape-based automatic differentiation. Common issues include watching non-trainable tensors, performing non-differentiable operations, or reading gradients outside tape scope. This guide shows practical patterns for gradient extraction in custom training logic.

Core Sections

1. Basic gradient computation

1import tensorflow as tf
2
3x = tf.Variable(3.0)
4
5with tf.GradientTape() as tape:
6    y = x * x + 2.0
7
8grad = tape.gradient(y, x)
9print(grad.numpy())  # 6.0

Tape records operations inside its context.

2. Gradients for model variables

1with tf.GradientTape() as tape:
2    preds = model(inputs, training=True)
3    loss = loss_fn(targets, preds)
4
5grads = tape.gradient(loss, model.trainable_variables)
6optimizer.apply_gradients(zip(grads, model.trainable_variables))

This is the core custom training-step pattern.

3. Watching non-variable tensors

1x = tf.constant(3.0)
2with tf.GradientTape() as tape:
3    tape.watch(x)
4    y = x ** 3
5print(tape.gradient(y, x))

Constants are not watched automatically.

4. Persistent tape for multiple gradients

1with tf.GradientTape(persistent=True) as tape:
2    y = f(x)
3    z = g(x)
4
5dy = tape.gradient(y, x)
6dz = tape.gradient(z, x)
7del tape

Release persistent tapes to avoid memory growth.

5. Gradient debugging

If gradients are None, check:

• variable is trainable/watched
• loss depends on variable path
• no tf.stop_gradient or non-diff ops break chain

6. model.fit integration

For most workflows, use model.compile + fit. Use custom tapes when you need custom losses, multi-optimizer logic, or gradient manipulation.

Validation and production readiness

A practical implementation should be validated beyond the happy path. Create a compact test matrix that includes standard input, boundary conditions, invalid data, and one realistic production-sized case. This reveals issues that unit-level examples often miss, such as silent coercions, ordering assumptions, and timeout behavior under load. If the workflow includes file or network operations, include at least one fault-injection test that simulates missing resources and transient failures.

1test_matrix:
2  - happy path: expected inputs and normal environment
3  - boundary path: min/max size, empty values, extreme ranges
4  - failure path: malformed input, unavailable dependency, timeout
5  - scale path: representative volume and concurrency

Operational safeguards are equally important. Add structured logging around the critical branches so you can diagnose failures quickly without reproducing them from scratch. A good log record should include operation name, key identifiers, and final outcome. Keep sensitive values masked. For asynchronous or background flows, include correlation IDs so related events can be traced across threads and services.

Define explicit fallback behavior before incidents occur. Decide whether the code should retry, fail fast, or degrade gracefully when dependencies are unavailable. If retries are used, bound them and use backoff. Unbounded retries often hide real outages and can amplify load problems. Add monitoring counters for success/failure/latency so regressions become visible immediately after deployment.

Finally, keep a short runbook near the code or documentation: required runtime versions, known platform differences, and a rollback plan. This turns one-off fixes into repeatable operational practices. Teams that standardize these checks usually reduce debugging time and avoid recurring reliability bugs.

Common Pitfalls

• Expecting gradients for constants without tape.watch.
• Computing loss outside tape scope.
• Ignoring None gradients and applying optimizer anyway.
• Keeping persistent tapes alive and leaking memory.
• Mixing NumPy operations in differentiable path.

Summary

In Keras/TensorFlow 2, tf.GradientTape is the standard way to extract gradients. Keep computations inside tape context, target trainable variables, and diagnose None gradients systematically. For standard training use fit; for advanced optimization logic use custom tape workflows.

Teams that document this exact approach in shared guidelines and enforce it through CI checks reduce repeated regressions, accelerate onboarding, and keep behavior consistent across local development, automated pipelines, and production operations.