Introduction

If you want a TensorFlow operation with a custom gradient and you want to stay in Python, the usual tool is @tf.custom_gradient. That is different from writing a compiled C++ custom op, which is more complex and only needed for low-level performance or unsupported kernels. For many research and model-development tasks, a Python-defined forward pass plus a custom backward rule is enough.

Use @tf.custom_gradient for Python-Level Custom Gradients

The decorator lets you define:

• the forward computation
• a function that returns gradients for the inputs

Here is a simple square function with an explicit gradient:

1import tensorflow as tf
2
3@tf.custom_gradient
4def square_with_grad(x):
5    y = x * x
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_with_grad(x)
16
17print(y.numpy())                     # 9.0
18print(tape.gradient(y, x).numpy())  # 6.0

This is the standard Python answer for “write an op with a gradient.”

Understand What the Gradient Function Receives

The grad function does not return the raw derivative alone. It receives dy, which is the gradient coming from later parts of the graph. Your job is to apply the chain rule and return gradients with respect to the input arguments.

For a single input:

returned gradient = dy * local_derivative

For multiple inputs, return one gradient per input in the same order the forward function accepted them.

Multi-Input Example

Here is a custom multiply operation with explicit gradients for both inputs.

1import tensorflow as tf
2
3@tf.custom_gradient
4def multiply_with_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
13
14a = tf.constant(4.0)
15b = tf.constant(5.0)
16
17with tf.GradientTape() as tape:
18    tape.watch([a, b])
19    out = multiply_with_grad(a, b)
20
21grads = tape.gradient(out, [a, b])
22print(out.numpy())       # 20.0
23print(grads[0].numpy())  # 5.0
24print(grads[1].numpy())  # 4.0

TensorFlow expects the gradient return structure to match the input structure.

Wrap Non-Differentiable or Special Logic Carefully

A common use case is expressing a forward rule that TensorFlow can execute, but whose default gradient is missing, unstable, or not what you want.

For example, you may want a clipped forward pass with a custom straight-through gradient:

1import tensorflow as tf
2
3@tf.custom_gradient
4def clipped_identity(x):
5    y = tf.clip_by_value(x, -1.0, 1.0)
6
7    def grad(dy):
8        return dy
9
10    return y, grad

This pattern is common in quantization research and surrogate-gradient methods.

Test the Gradient Instead of Assuming It Is Correct

A custom gradient can silently be wrong. Always test it against expectations.

1import tensorflow as tf
2
3x = tf.constant(2.0)
4
5with tf.GradientTape() as tape:
6    tape.watch(x)
7    y = square_with_grad(x)
8
9grad = tape.gradient(y, x)
10tf.debugging.assert_near(grad, tf.constant(4.0))
11print("gradient check passed")

For more complex functions, compare against finite-difference approximations or a reference implementation.

Know When Python Is Not Enough

@tf.custom_gradient is great for gradient logic, but it does not create a brand-new low-level TensorFlow kernel. If you need:

• custom device kernels
• low-level op registration
• optimized C++ or CUDA performance

then you are in true custom-op territory, which usually means compiled extensions, not just Python code.

That distinction matters because many questions use the word “op” loosely. In a lot of cases, what they really need is a Python-defined differentiable function, not a compiled TensorFlow op.

Use with tf.function and Keras

Python-defined custom gradients usually work fine inside tf.function and Keras models.

1import tensorflow as tf
2
3inputs = tf.keras.Input(shape=(1,))
4outputs = square_with_grad(inputs)
5model = tf.keras.Model(inputs, outputs)
6
7result = model(tf.constant([[3.0]]))
8print(result.numpy())

That makes @tf.custom_gradient practical for real model code, not just isolated experiments.

Common Pitfalls

One common mistake is returning the raw derivative instead of multiplying by the incoming dy.

Another issue is forgetting to return one gradient per input argument in multi-input functions.

A third mistake is assuming @tf.custom_gradient creates a new low-level compiled TensorFlow op. It does not; it customizes autodiff behavior around Python-defined logic.

Summary

• Use @tf.custom_gradient when you want a custom differentiable operation in Python.
• Define both the forward value and the backward gradient rule explicitly.
• Return gradients in the same structure as the input arguments.
• Test custom gradients instead of assuming the math is correct.
• Move to compiled custom ops only when Python-level customization is not enough.