Introduction

If an operation is built from ordinary TensorFlow ops, you usually do not need to register a low-level gradient override in the graph. The modern and simplest approach is to wrap the forward computation with @tf.custom_gradient and return a backward function that defines how gradients should propagate.

Why Use a Custom Gradient

TensorFlow's automatic differentiation is often enough, but there are cases where you want to define the backward pass yourself:

• the default gradient is numerically unstable
• the backward formula can be made cheaper than the automatic one
• the operation is mathematically correct but you want a special training-time surrogate gradient
• you want to stop or reshape gradient flow deliberately

When the operation is composed of TensorFlow ops rather than a custom C++ kernel, @tf.custom_gradient is usually the right abstraction.

A Minimal @tf.custom_gradient Example

The decorated function returns the forward result and a nested gradient function.

1import tensorflow as tf
2
3@tf.custom_gradient
4def square_with_scaled_grad(x):
5    y = x * x
6
7    def grad(dy):
8        return dy * 3.0 * x
9
10    return y, grad
11
12x = tf.Variable(2.0)
13with tf.GradientTape() as tape:
14    y = square_with_scaled_grad(x)
15
16print(tape.gradient(y, x).numpy())  # 6.0

The forward pass still computes x * x, but the backward pass is now whatever grad returns.

Apply It to a Composite TensorFlow Operation

Suppose your forward computation is built from several TensorFlow ops and you want to keep the forward logic readable while defining a custom backward rule.

1import tensorflow as tf
2
3@tf.custom_gradient
4def stable_log1pexp(x):
5    y = tf.math.log1p(tf.exp(-tf.abs(x))) + tf.maximum(x, 0.0)
6
7    def grad(dy):
8        return dy * tf.math.sigmoid(x)
9
10    return y, grad
11
12x = tf.constant([-2.0, 0.0, 3.0])
13with tf.GradientTape() as tape:
14    tape.watch(x)
15    y = stable_log1pexp(x)
16
17print(tape.gradient(y, x).numpy())

This is a good example of a composite op whose gradient you may want to control explicitly for stability or clarity.

Multiple Inputs Are Supported

If the custom operation takes several inputs, the gradient function must return one gradient per input.

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

The order matters. The returned gradients must line up with the order of the original inputs.

Use It with GradientTape

The easiest way to test a custom gradient is with tf.GradientTape. That gives immediate confirmation that the backward rule is being used.

1x = tf.Variable(3.0)
2with tf.GradientTape() as tape:
3    y = square_with_scaled_grad(x)
4
5print(float(tape.gradient(y, x)))

If the result does not match your intended derivative, fix the gradient function before embedding the op in a larger model.

When to Reach for Lower-Level Gradient Registration

Older TensorFlow 1 style code sometimes used graph-level gradient overrides and registration hooks. Those still exist in some legacy workflows, but they are usually not necessary when the operation is just a Python function composed of TensorFlow ops.

The practical rule is:

• composite Python op: prefer @tf.custom_gradient
• true custom kernel or lower-level graph integration: use the lower-level registration tools if needed

For most TensorFlow 2 code, @tf.custom_gradient is the clean answer.

Be Careful with Correctness

A custom gradient can intentionally differ from the mathematical derivative, but if you do that, it should be a deliberate modeling decision. Otherwise you can silently train a model on incorrect gradients.

That is why custom-gradient code should be small, tested, and well named. It is one of the easiest places to create a model that runs fine and still learns the wrong thing.

Common Pitfalls

• Reaching for low-level registration when @tf.custom_gradient is enough.
• Returning the wrong number of gradients for multi-input functions.
• Writing a custom backward rule without testing it under GradientTape.
• Forgetting that the gradient function receives upstream gradient dy.
• Using a surrogate gradient without documenting that it intentionally differs from the true derivative.

Summary

• For a TensorFlow operation composed of TensorFlow ops, @tf.custom_gradient is usually the right tool.
• Return the forward result plus a nested function that defines the backward pass.
• The gradient function must account for upstream gradient flow through dy.
• Multi-input functions must return one gradient per input.
• Test the custom backward rule explicitly before using it in model training.