Introduction

The grad_ys parameter in tf.gradients(ys, xs, grad_ys) is one of the most overlooked yet powerful features in TensorFlow's automatic differentiation system. It lets you control the initial gradient that flows backward through the computation graph. Understanding what it does and why you would use it is essential for implementing custom training loops, multi-task learning, and advanced gradient manipulation techniques.

What tf.gradients Does

tf.gradients(ys, xs) computes the symbolic partial derivatives of the tensors in ys with respect to the tensors in xs. Internally, TensorFlow walks the computation graph backward from ys to xs, applying the chain rule at each operation node.

1import tensorflow as tf
2
3# TensorFlow 1.x style
4x = tf.Variable(3.0)
5y = x ** 2  # y = x^2
6
7grads = tf.gradients(y, x)  # dy/dx = 2x = 6.0

The function returns a list of tensors, one for each element in xs, where each tensor is the sum of all partial derivatives of ys with respect to that xs element.

What grad_ys Represents

The grad_ys parameter is the initial gradient, sometimes called the upstream gradient or the seed gradient. In the chain rule, when computing dL/dx where L is a downstream loss, you are really computing (dL/dy) * (dy/dx). The grad_ys parameter is that dL/dy term.

When you omit grad_ys (or pass None), TensorFlow defaults it to a tensor of ones with the same shape as ys. This means tf.gradients(y, x) computes dy/dx as if y were the final loss scalar, with an implicit upstream gradient of 1.0.

1import tensorflow as tf
2
3x = tf.Variable(3.0)
4y = x ** 2
5
6# These two calls are equivalent
7grads_default = tf.gradients(y, x)               # grad_ys defaults to [1.0]
8grads_explicit = tf.gradients(y, x, grad_ys=[1.0])  # Same result

Practical Use Cases

Weighted Gradients for Multi-Task Learning

When you have multiple loss functions, grad_ys lets you weight their contributions to the final gradient without modifying the loss computation itself:

1import tensorflow as tf
2
3x = tf.Variable(2.0)
4loss_a = x ** 2      # Task A loss
5loss_b = (x - 5) ** 2  # Task B loss
6
7# Weight task A at 0.3 and task B at 0.7
8grads = tf.gradients(
9    [loss_a, loss_b],
10    x,
11    grad_ys=[0.3, 0.7]
12)
13# Result: 0.3 * d(loss_a)/dx + 0.7 * d(loss_b)/dx
14# = 0.3 * (2*2) + 0.7 * (2*(2-5))
15# = 0.3 * 4 + 0.7 * (-6) = 1.2 - 4.2 = -3.0

Gradient Scaling

You can scale all gradients by a constant factor, which is useful for gradient accumulation across micro-batches:

scale = 1.0 / num_micro_batches
grads = tf.gradients(loss, weights, grad_ys=[scale])

Stopping Gradient Flow Selectively

By setting certain elements of grad_ys to zero, you can block gradient flow from specific outputs:

1import tensorflow as tf
2
3x = tf.Variable(1.0)
4y1 = x * 3
5y2 = x * 5
6
7# Only compute gradient from y1, ignore y2
8grads = tf.gradients([y1, y2], x, grad_ys=[1.0, 0.0])
9# Result: 3.0 (only dy1/dx)

TF2 Equivalent with GradientTape

In TensorFlow 2.x, tf.gradients is replaced by tf.GradientTape for eager execution. The equivalent of grad_ys is the output_gradients parameter in the gradient() method:

1import tensorflow as tf
2
3x = tf.Variable(3.0)
4
5with tf.GradientTape() as tape:
6    y = x ** 2
7
8# Default gradient (output_gradients=1.0 implicitly)
9grad_default = tape.gradient(y, x)
10print(grad_default)  # 6.0
11
12# With explicit output_gradients (equivalent to grad_ys)
13with tf.GradientTape() as tape:
14    y = x ** 2
15
16grad_scaled = tape.gradient(y, x, output_gradients=2.0)
17print(grad_scaled)  # 12.0 (= 2.0 * 2*3)

For multiple outputs with different weights:

1import tensorflow as tf
2
3x = tf.Variable(2.0)
4
5with tf.GradientTape() as tape:
6    y1 = x ** 2
7    y2 = x ** 3
8
9grad = tape.gradient(
10    [y1, y2],
11    x,
12    output_gradients=[0.5, 0.5]
13)
14# 0.5 * dy1/dx + 0.5 * dy2/dx
15# = 0.5 * 4 + 0.5 * 12 = 8.0
16print(grad)  # 8.0

Relationship to the Chain Rule

The mathematical foundation is straightforward. Given a computation chain L -> y -> x, the chain rule states:

dL/dx = (dL/dy) * (dy/dx)

In tf.gradients(y, x, grad_ys):

• y and x define the (dy/dx) portion
• grad_ys provides the (dL/dy) portion
• The return value is the full (dL/dx)

When grad_ys is None, TensorFlow substitutes dL/dy = 1, which is correct when y is the final loss and there is no further upstream computation.

Common Pitfalls

• Forgetting that grad_ys defaults to ones, not zeros, which means omitting it computes the full gradient rather than no gradient.
• Passing a grad_ys tensor with a shape that does not match ys, which causes a shape mismatch error.
• Confusing grad_ys with gradient clipping; grad_ys scales the initial seed, not the final computed gradient.
• Using tf.gradients inside eager mode in TF2 without wrapping it in tf.function, since it operates on symbolic graphs.
• Not realizing that when ys is a list, grad_ys must also be a list of the same length.

Summary

• grad_ys is the upstream (initial) gradient seed passed into the backward pass of tf.gradients.
• When omitted, it defaults to a tensor of ones, meaning "compute dy/dx directly."
• Use it for weighted gradients in multi-task learning, gradient scaling, and selective gradient blocking.
• In TensorFlow 2.x, the equivalent is the output_gradients parameter in GradientTape.gradient().
• Always ensure grad_ys has the same shape and length as ys to avoid runtime errors.