TensorFlow
Gradient Calculation
Machine Learning
Neural Networks
Debugging

Is gradient in the tensorflow's graph calculated incorrectly?

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design

Introduction

TensorFlow gradients are usually not "wrong" in the general sense. TensorFlow uses automatic differentiation, so for standard differentiable operations the gradients follow the chain rule exactly as defined by the registered gradient implementations. When developers think the graph is computing gradients incorrectly, the real issue is usually shape reduction, disconnected tensors, nondifferentiable operations, or a misunderstanding of what GradientTape.gradient returns.

How TensorFlow Computes Gradients

In TensorFlow, operations executed under tf.GradientTape are recorded. When you ask for a gradient, TensorFlow walks backward through the recorded operations and applies gradient rules.

python
1import tensorflow as tf
2
3x = tf.Variable(3.0)
4
5with tf.GradientTape() as tape:
6    y = x * x + 2.0 * x
7
8grad = tape.gradient(y, x)
9print(grad.numpy())

The derivative of x^2 + 2x at x = 3 is 8, and TensorFlow returns that value.

So when the result looks surprising, start by asking whether the mathematical object you requested is the one you intended.

One Common Surprise: Gradients of a Sum

If the target tensor is not scalar, tape.gradient computes the gradient of the sum of that target unless you explicitly ask for Jacobian-like behavior.

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

This prints 7.0, not [3.0, 4.0], because TensorFlow differentiates the sum of the target elements. That behavior is documented and often misread as a bug.

If you need elementwise derivatives, use tape.jacobian instead.

None Does Not Mean Bad Math

Another common confusion is getting None as the gradient. That usually means the source tensor is not connected to the target, is not being watched, or uses a nondifferentiable dtype such as integer tensors.

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

This returns None because integer tensors do not participate in standard floating-point differentiation the way trainable float variables do.

Debugging with a Numerical Check

If you truly suspect a bad gradient, compare TensorFlow's analytic gradient with a finite-difference approximation.

python
1import tensorflow as tf
2
3x = tf.Variable(1.5)
4
5with tf.GradientTape() as tape:
6    y = tf.math.sin(x) * x
7analytic = tape.gradient(y, x).numpy()
8
9eps = 1e-4
10f_plus = (tf.math.sin(x + eps) * (x + eps)).numpy()
11f_minus = (tf.math.sin(x - eps) * (x - eps)).numpy()
12numeric = (f_plus - f_minus) / (2 * eps)
13
14print("analytic:", analytic)
15print("numeric:", numeric)

The numbers should be close. If they are not, look for one of these causes first:

  • the function is numerically unstable near that point
  • the epsilon value is poorly chosen
  • the graph contains custom operations or custom gradients
  • the tensor path is disconnected or stateful in a way you did not expect

Cases That Really Can Mislead You

There are real edge cases, but they are usually not TensorFlow miscomputing a standard derivative. Typical examples include:

  • nondifferentiable ops such as hard thresholding or argmax-like logic
  • custom layers with wrong @tf.custom_gradient logic
  • stateful updates that break the differentiation path
  • gradients through values that were converted to NumPy or Python scalars

In those cases, the problem is normally in the modeled computation, not in the core gradient engine.

Use the Right API for the Question

If you need:

  • one gradient for a scalar loss, use tape.gradient
  • per-output derivatives, use tape.jacobian
  • custom derivative logic, use @tf.custom_gradient

A lot of "incorrect gradient" reports come from using the scalar-loss API while expecting full Jacobian behavior.

Common Pitfalls

The biggest mistake is asking for the gradient of a vector target and expecting TensorFlow to return per-element derivatives automatically.

Another mistake is differentiating through integer tensors, Python numbers, or values that left the TensorFlow graph.

A third issue is assuming None means a wrong derivative instead of an unconnected or unwatched path.

Summary

  • TensorFlow gradients for standard differentiable ops are usually correct
  • 'GradientTape.gradient returns the gradient of the summed target when the target is non-scalar'
  • 'None usually means disconnected, unwatched, or nondifferentiable sources'
  • Use numerical finite differences when you need to sanity-check a suspicious result
  • Reach for jacobian or custom-gradient tools when the default scalar-loss API is not the right abstraction

Related reading
Free course
Beginner
7 lessons
2 hours
Tackling System Design Interview Problems

A short course that equips you with the skills to approach system design interviews methodically.

Start the free course
Track what you have practised

A free account saves your progress, solutions and study plan across every problem on Codemia.

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design

All Rights Reserved.