Introduction

In TensorFlow 2, the normal way to create a Keras layer with a custom gradient is to keep the layer as a regular tf.keras.layers.Layer subclass and place the nonstandard backward rule inside a function decorated with @tf.custom_gradient. The layer handles weights and shape logic, while the decorated function defines how backpropagation should treat the operation.

Separate The Layer From The Gradient Rule

A common mistake is trying to override backpropagation by subclassing Layer alone. That is not where TensorFlow expects gradient logic. The custom backward rule belongs in an operation, and the layer simply calls that operation.

That structure looks like this:

• define an op with @tf.custom_gradient
• subclass Layer
• call the op from call
• use the layer in an ordinary Keras model

A Minimal Example

The example below implements a layer that behaves like squaring in the forward pass, but uses a clipped gradient in the backward pass.

1import tensorflow as tf
2
3@tf.custom_gradient
4def square_with_clipped_grad(x):
5    y = tf.square(x)
6
7    def grad(dy):
8        raw_grad = 2.0 * x
9        clipped_grad = tf.clip_by_value(raw_grad, -1.0, 1.0)
10        return dy * clipped_grad
11
12    return y, grad
13
14class CustomSquareLayer(tf.keras.layers.Layer):
15    def call(self, inputs):
16        return square_with_clipped_grad(inputs)
17
18x = tf.constant([[2.0], [3.0]])
19layer = CustomSquareLayer()
20print(layer(x))

The forward output is the same as tf.square. The difference appears during training when gradients flow backward.

Put The Layer Inside A Model

Because the custom op is wrapped inside a normal layer, you can use it in a standard Keras model.

1import tensorflow as tf
2
3inputs = tf.keras.Input(shape=(1,))
4x = CustomSquareLayer()(inputs)
5outputs = tf.keras.layers.Dense(1)(x)
6model = tf.keras.Model(inputs, outputs)
7
8model.compile(optimizer="adam", loss="mse")
9
10x_train = tf.constant([[1.0], [2.0], [3.0], [4.0]])
11y_train = tf.constant([[1.0], [4.0], [9.0], [16.0]])
12
13model.fit(x_train, y_train, epochs=2, verbose=0)

From Keras's point of view, this is just another differentiable layer in the graph.

When A Custom Gradient Is Appropriate

You should only define a custom gradient when the default automatic differentiation rule is not what you want.

Common reasons are:

• the forward op is nondifferentiable but you want a surrogate gradient
• the mathematically correct gradient is unstable and you want a numerically safer one
• you are implementing a research idea such as a straight-through estimator
• you want to clip, rescale, or reshape the backward signal deliberately

If TensorFlow's built-in gradient is already correct and stable, do not replace it just to make the layer feel more advanced.

A Straight-Through Estimator Style Example

A common use case is discretization in the forward pass with a smoother gradient in the backward pass.

1import tensorflow as tf
2
3@tf.custom_gradient
4def round_with_identity_grad(x):
5    y = tf.round(x)
6
7    def grad(dy):
8        return dy
9
10    return y, grad
11
12class QuantizeLayer(tf.keras.layers.Layer):
13    def call(self, inputs):
14        return round_with_identity_grad(inputs)

This style is not mathematically exact, but it is sometimes useful in quantization-aware or discrete-latent experiments.

Gradients For Variables

If your custom op closes over trainable variables, the gradient function may also need to return gradients for those variables. In many layer designs, though, the custom gradient only depends on the tensor input and the layer's trainable weights are handled by ordinary TensorFlow ops around it.

The simpler your custom op boundary is, the easier it is to debug.

Debugging The Gradient

When you introduce a custom gradient, validate it explicitly. GradientTape is the fastest way to do that.

1import tensorflow as tf
2
3x = tf.Variable([[2.0]])
4
5with tf.GradientTape() as tape:
6    y = square_with_clipped_grad(x)
7
8print(tape.gradient(y, x))

If the gradient is not what you expect here, the model training loop will not save you later.

Keep Serialization In Mind

If you want the layer to be saved and reloaded cleanly as part of a Keras model, implement get_config when the layer has constructor parameters.

1class CustomSquareLayer(tf.keras.layers.Layer):
2    def __init__(self, **kwargs):
3        super().__init__(**kwargs)
4
5    def call(self, inputs):
6        return square_with_clipped_grad(inputs)
7
8    def get_config(self):
9        return super().get_config()

That is not specific to custom gradients, but custom layers often get reused in saved models, so it matters.

Common Pitfalls

• Trying to put gradient logic directly into the layer class instead of a @tf.custom_gradient function.
• Replacing a correct built-in gradient without a clear reason.
• Forgetting to validate the backward rule with GradientTape.
• Writing a custom gradient that has the wrong shape or dtype relative to the inputs.
• Building a custom op that depends on hidden side effects rather than pure tensor logic.

Summary

• In TF2, define the custom backward rule with @tf.custom_gradient and call it from a Keras layer.
• The layer is responsible for model integration; the decorated function is responsible for gradient behavior.
• This approach works naturally inside standard Keras models and training loops.
• Use custom gradients only when the default autodiff rule is not the behavior you want.
• Validate the gradient explicitly before trusting model training.