Introduction

Custom gradients in TF2/Keras are useful when default autodiff is unavailable, numerically unstable, or intentionally overridden (for example straight-through estimators, clipping tricks, or custom surrogate gradients). In Keras layers, the clean pattern is to wrap the custom operation with @tf.custom_gradient and call it inside Layer.call.

This keeps integration with Model.fit intact while giving precise control over backward pass behavior.

Core Sections

1. Define a custom op with gradient

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        # default would be 2*x; here we scale it intentionally
9        return dy * (0.5 * 2.0 * x)
10
11    return y, grad

The function returns forward output and a gradient function.

2. Wrap it in a Keras layer

class CustomSquareLayer(tf.keras.layers.Layer):
    def call(self, inputs):
        return square_with_scaled_grad(inputs)

Now layer can be used in normal models.

3. Build and train model as usual

1model = tf.keras.Sequential([
2    tf.keras.layers.Input(shape=(1,)),
3    CustomSquareLayer(),
4    tf.keras.layers.Dense(1)
5])
6
7model.compile(optimizer="adam", loss="mse")
8model.fit(x_train, y_train, epochs=3)

Keras will use your custom backward rule during optimization.

4. Verify gradients explicitly

1x = tf.Variable([[2.0]])
2with tf.GradientTape() as tape:
3    y = square_with_scaled_grad(x)
4
5print(tape.gradient(y, x).numpy())

Always test custom gradients before large training runs.

5. Handle non-differentiable regions carefully

If gradient is undefined at certain points, choose and document surrogate behavior explicitly. Silent assumptions can destabilize training.

Common Pitfalls

• Implementing custom gradients without validating with GradientTape tests.
• Returning wrong gradient shape or dtype from grad function.
• Overriding gradients without clear mathematical justification.
• Mixing NumPy operations inside custom-gradient functions and breaking graph compatibility.
• Ignoring edge cases around non-differentiable regions or clipping boundaries.

Summary

In TF2, create Keras layers with custom gradients by defining @tf.custom_gradient functions and calling them inside layer call. This gives full control over backward behavior while preserving normal model training APIs. Validate gradient correctness early, document intentional deviations, and keep gradient implementations graph-safe for stable learning.

A practical way to make this guidance durable is to convert it into a small runbook that includes prerequisites, expected environment versions, and a short verification sequence. Even strong teams lose time when troubleshooting steps live only in memory or chat history. A runbook should explicitly answer three questions: what to check first, what output confirms healthy behavior, and what output indicates a known failure mode. This level of clarity helps both experienced maintainers and newer contributors, and it reduces repeated investigation during incidents.

It is also valuable to create a tiny reproducible fixture for this topic. The fixture can be a minimal script, test case, sample request, or small dataset that demonstrates the correct behavior in isolation. When regressions appear after dependency upgrades, infrastructure changes, or framework migrations, that fixture becomes the fastest way to isolate whether the issue is environmental or logic-related. Keeping a focused fixture in source control gives you a stable benchmark across branches and release cycles.

For long-term reliability, pair documentation with one automated guardrail in CI. The guardrail should be narrow and fast: an import check, schema validation, endpoint contract test, deterministic unit test, or lightweight performance threshold. Avoid broad flaky checks that hide real signals. The goal is early, actionable feedback before code reaches production. If the same category of issue appears repeatedly, promote the manual troubleshooting step into automation so the system catches it first. Over time, this shifts effort from reactive debugging to preventive quality control and keeps the knowledge article relevant in real engineering workflows.

As a final hardening step, periodically rerun the sample code in a clean environment image and record results in version control. This catches ecosystem drift early and keeps implementation guidance aligned with real runtime behavior.