Introduction

If you want TensorFlow to keep the forward computation of a subgraph but use custom backpropagation through it, the standard tool is a custom gradient. In modern TensorFlow, the cleanest approach is to wrap the relevant subgraph in a function decorated with tf.custom_gradient so only that section gets overridden while the rest of the model still uses normal automatic differentiation.

Core Sections

Why define a gradient for only part of the graph

There are several valid reasons to override gradients selectively:

• the forward operation is numerically fine but the default gradient is unstable
• you want a straight-through estimator for a non-differentiable step
• you want to clip, scale, or reshape gradients through one subgraph only
• you are wrapping an operation whose analytic derivative you know explicitly

The important point is scope. You usually do not want to rewrite the gradient rules for the whole model, only for one problematic transformation.

Use tf.custom_gradient around the target subgraph

In TensorFlow 2, define a function whose body performs the normal forward computation and whose nested gradient function returns the custom derivative.

1import tensorflow as tf
2
3@tf.custom_gradient
4def stable_square(x):
5    y = x * x
6
7    def grad(dy):
8        return dy * tf.clip_by_value(2.0 * x, -10.0, 10.0)
9
10    return y, grad
11
12x = tf.constant([3.0, -7.0])
13with tf.GradientTape() as tape:
14    tape.watch(x)
15    loss = tf.reduce_sum(stable_square(x))
16
17print(tape.gradient(loss, x))

Only stable_square uses the custom gradient. The rest of the computation graph continues to use TensorFlow’s normal gradient rules.

Think of the wrapper as the subgraph boundary

The “subgraph” here is whatever you put inside the wrapped function. If you need a larger custom region, move more operations inside. If you only need to override one step, keep the wrapper small. Smaller boundaries are easier to test and reason about.

For example, you could wrap a binarization step while leaving the rest of the model untouched.

1@tf.custom_gradient
2def straight_through_round(x):
3    y = tf.round(x)
4
5    def grad(dy):
6        return dy
7
8    return y, grad

This is a classic straight-through estimator: the forward pass rounds, but the backward pass pretends the operation was identity.

Test the custom gradient in isolation

Before inserting it into a large training graph, check the local behavior explicitly with GradientTape. That helps you catch sign errors, shape mismatches, or unintentionally exploding derivatives early.

Custom gradients are powerful, but they can also silently make optimization worse if the derivative you return does not match the intended training behavior.

Older graph-mode alternatives

In older TensorFlow graph code, people sometimes used gradient overrides through graph-level mechanisms. Those still exist in some legacy workflows, but they are heavier and easier to misuse. For modern eager TensorFlow and Keras code, tf.custom_gradient is usually the right answer.

If you are working with a custom C++ op or a legacy TensorFlow 1 graph, there are lower-level registration mechanisms. But for Python-defined subgraphs, the decorator is the most direct path.

Common Pitfalls

• Overriding a large region of the graph when only one unstable or non-differentiable step actually needs a custom gradient.
• Returning gradients with the wrong shape or dtype for the wrapped function inputs.
• Treating the custom gradient as a mathematical formality and never testing how it behaves under GradientTape before training.
• Using legacy graph override techniques in modern TensorFlow when tf.custom_gradient would be simpler and clearer.
• Forgetting that the custom derivative changes optimization behavior even if the forward pass still looks identical.

Summary

• To define the gradient for only a TensorFlow subgraph, wrap that subgraph in a function decorated with tf.custom_gradient.
• The forward pass stays local to the wrapper, and the nested gradient function controls backpropagation through that region.
• Keep the wrapped region as small as possible so the custom behavior is easy to reason about.
• Test custom derivatives explicitly before integrating them into a full model.
• In modern TensorFlow, tf.custom_gradient is the standard solution for this problem.