Gradient computation

Automatic differentiation. Two interfaces, one shared registry.

The two interfaces

Eager: tf.GradientTape

The modern API. A context manager that records every op while it's open:

x = tf.Variable(3.0)
with tf.GradientTape() as tape:
    y = x ** 2 + 2 * x + 1
dy_dx = tape.gradient(y, x)   # 2*x + 2 = 8

Implementation:

• Python: tensorflow/python/eager/tape.py, tensorflow/python/eager/backprop.py.
• C++ (the actual recorder): tensorflow/c/eager/tape.h and the wrappers in tensorflow/python/eager/.

Tapes can be persistent (read multiple times), can watch non-Variable tensors, and can be nested for higher-order derivatives.

Graph mode: tf.gradients

The older API for graph mode (tf.compat.v1.gradients). Still used internally and in some legacy code paths. Implementation in tensorflow/python/ops/gradients_impl.py and gradients_util.py.

Both APIs lower to the same gradient-op registry.

The gradient registry

Every op that participates in autodiff has a gradient registered. The decorator @RegisterGradient("OpName") in tensorflow/python/framework/ops.py takes a function (op, grad) -> [grad_in0, grad_in1, ...]:

@ops.RegisterGradient("Add")
def _AddGrad(op, grad):
    x = op.inputs[0]
    y = op.inputs[1]
    sx = array_ops.shape(x)
    sy = array_ops.shape(y)
    rx, ry = gen_array_ops.broadcast_gradient_args(sx, sy)
    return (array_ops.reshape(math_ops.reduce_sum(grad, rx), sx),
            array_ops.reshape(math_ops.reduce_sum(grad, ry), sy))

Gradient implementations live in tensorflow/python/ops/<family>_grad.py:

File	Covers
tensorflow/python/ops/math_grad.py	Add, Sub, Mul, MatMul, ...
tensorflow/python/ops/array_grad.py	Reshape, Slice, Pad, ...
tensorflow/python/ops/nn_grad.py	Relu, Softmax, Conv2D, ...
tensorflow/python/ops/control_flow_grad.py	Switch, Merge, while-loop gradients
tensorflow/python/ops/data_flow_grad.py	Queues, scatter/gather
tensorflow/python/ops/sparse_grad.py	Sparse tensor ops
tensorflow/python/ops/linalg_grad.py	Linear algebra
tensorflow/python/ops/state_grad.py	Variable-related ops

C++ gradient registry

For the C++ frontend, gradients live in tensorflow/cc/gradients/:

REGISTER_GRADIENT_OP("Add", AddGrad);

Used by tensorflow/cc/framework/gradients.cc to compute symbolic gradients in C++ (most users come in through Python and don't touch this path).

How GradientTape records

sequenceDiagram
    participant Py as Python op call
    participant Tape as Active GradientTape
    participant Pybind as eager runtime
    participant Cap as RecordOperation

    Py->>Tape: any active tape?
    alt yes
        Tape->>Cap: capture (op_name, inputs, outputs, attrs, backward_function)
    end
    Py->>Pybind: execute the op

When a tape is open, record_operation (C++) appends a node to a per-thread stack. Each node carries the op's inputs, outputs, and a backward function — typically a Python lambda that calls into the registered gradient function. On tape.gradient(target, sources), the tape walks backwards, multiplying gradients along the way.

For ops not registered as differentiable, tape.gradient returns None for that source.

jvp and forward-mode

TensorFlow has experimental forward-mode autodiff: tf.autodiff.ForwardAccumulator (tensorflow/python/eager/forwardprop.py). Less heavily used than reverse mode (GradientTape) but supported for higher-order derivatives.

Variable gradients

Most users care about gradients with respect to tf.Variables. The tape automatically watches all variables created or accessed inside its scope; tape.watch(t) is needed only for non-variable tensors.

Composite gradients

Some ops have structural gradients — for example tf.cond and tf.while_loop's gradient ops live in tensorflow/python/ops/control_flow_grad.py and reconstruct the body's gradient by calling tf.gradients on the body's traced graph. This is one of the more complex pieces of the gradient system.

Stop gradient and custom gradients

• tf.stop_gradient(x) blocks gradient flow through x. Implemented as the StopGradient op with a registered gradient that returns None.
• tf.custom_gradient lets you override the gradient of an arbitrary computation — define a forward and backward function and TF takes care of the rest.

Integration points

• tf.function — variables captured into a traced function still flow gradients correctly.
• Distributed — gradients on MirroredVariables are automatically all-reduced before the optimizer applies them.
• XLA — jit_compile=True recovers gradients from the traced function before lowering, so XLA sees a fused forward+backward HLO module.

Entry points for modification

• New gradient: write a @RegisterGradient("Op") function in the appropriate *_grad.py file. Tests live in tensorflow/python/kernel_tests/.
• New tape feature: tensorflow/python/eager/tape.py, backprop.py, and the C++ tape in tensorflow/c/eager/tape.h.
• New control-flow gradient: tensorflow/python/ops/control_flow_grad.py.

Related

• systems/eager-execution — the runtime layer the tape rides on.
• systems/kernels-and-ops — registered ops; gradients are registered against op names.
• features/tf-function-and-autograph — gradients inside @tf.function.