Introduction

If an operation is built from normal TensorFlow ops, the best way to give it a custom backward pass in modern TensorFlow is usually @tf.custom_gradient. That is different from the older graph-level gradient registration APIs, which are mainly for overriding primitive ops in TensorFlow 1 style graphs. For a composed Python function, the clean solution is to wrap the forward computation and explicitly return the gradient function.

Use @tf.custom_gradient for Composite Operations

Suppose you have a forward operation composed of ordinary TensorFlow ops and you want a custom derivative for numerical stability or modeling reasons.

1import tensorflow as tf
2
3@tf.custom_gradient
4def square_plus_one(x):
5    y = tf.square(x) + 1.0
6
7    def grad(dy):
8        return dy * 2.0 * x
9
10    return y, grad
11
12x = tf.constant(3.0)
13with tf.GradientTape() as tape:
14    tape.watch(x)
15    y = square_plus_one(x)
16
17print(tape.gradient(y, x).numpy())

The forward pass uses standard TensorFlow ops. The custom part is the nested grad function, which receives the upstream gradient dy and returns the downstream gradient for the input.

Understand the Contract of the Gradient Function

The gradient function must return one gradient per differentiable input. If your forward function takes multiple tensors, the custom gradient must return the same number of corresponding gradients.

1import tensorflow as tf
2
3@tf.custom_gradient
4def multiply_with_custom_grad(x, y):
5    z = x * y
6
7    def grad(dz):
8        dx = dz * y
9        dy = dz * x
10        return dx, dy
11
12    return z, grad

This explicit structure is why @tf.custom_gradient is so useful. It keeps the custom derivative tied directly to the function being defined.

Save Intermediate Values Deliberately

A custom gradient can close over tensors from the forward pass. That is often cleaner than recomputing everything inside the gradient function.

1import tensorflow as tf
2
3@tf.custom_gradient
4def stabilized_log1pexp(x):
5    y = tf.math.log1p(tf.exp(-tf.abs(x))) + tf.maximum(x, 0.0)
6
7    def grad(dy):
8        sigmoid = tf.math.sigmoid(x)
9        return dy * sigmoid
10
11    return y, grad

This pattern is common when the default symbolic gradient is correct mathematically but less stable numerically than a hand-written derivative.

Know When You Do Not Need a Custom Gradient

If your operation is already composed of differentiable TensorFlow ops and you are happy with the default derivative, you do not need to register anything. TensorFlow can differentiate through the graph automatically.

Use a custom gradient only when you actually need one of these:

1. A more stable derivative.
2. A simplified or clipped backward pass.
3. A straight-through estimator.
4. A domain-specific gradient that intentionally differs from the default symbolic derivative.

If the goal is only to wrap several TensorFlow ops into a helper function, the default autodiff is usually enough.

Distinguish This from tf.RegisterGradient

Older TensorFlow examples often mention tf.RegisterGradient. That API is graph-oriented and is mainly used when overriding gradients for named primitive ops in TensorFlow 1 style execution. For a Python function composed from existing ops, @tf.custom_gradient is normally the better fit.

A practical rule:

1. Composite Python function: use @tf.custom_gradient.
2. Low-level graph override for a specific op type: legacy registration APIs.

Keeping that distinction clear avoids a lot of unnecessary graph-manipulation complexity.

Test the Gradient Explicitly

Whenever you write a custom gradient, test it. A forward function that appears correct can still return the wrong backward values if the gradient function is incomplete or shaped incorrectly.

1import tensorflow as tf
2
3x = tf.constant(2.0)
4with tf.GradientTape() as tape:
5    tape.watch(x)
6    y = square_plus_one(x)
7
8grad = tape.gradient(y, x)
9print(grad.numpy())

For more complex functions, compare the result to the known analytical derivative or to a numerical approximation.

Common Pitfalls

• Reaching for legacy gradient-registration APIs when the operation is really just a composite Python function.
• Returning the wrong number of gradients from the custom gradient function.
• Forgetting that the gradient function receives the upstream gradient dy and must multiply through it appropriately.
• Writing a custom gradient where TensorFlow's default autodiff already does the right thing.
• Failing to test the backward pass explicitly after defining the forward function.

Summary

• For an operation composed of normal TensorFlow ops, @tf.custom_gradient is usually the correct modern tool.
• The forward function returns both the output and a nested gradient function.
• The gradient function must return one gradient per differentiable input.
• Custom gradients are useful for stability, modeling tricks, and deliberate backward-pass control.
• If the default TensorFlow derivative is already correct and sufficient, no registration is needed.