Introduction

The Jacobian is the matrix of first-order partial derivatives for a vector-valued function. In modern TensorFlow, the standard way to compute it is tf.GradientTape().jacobian(...), which gives you derivatives of every output component with respect to every input component.

Jacobian Versus Gradient

A gradient is for a scalar output. A Jacobian is for a vector or higher-rank output.

If y = f(x) and y has shape (m,) while x has shape (n,), the Jacobian has shape (m, n). TensorFlow generalizes that rule: the output tensor shape comes first, and the source tensor shape comes after it.

Basic TensorFlow Example

1import tensorflow as tf
2
3x = tf.Variable([1.0, 2.0], dtype=tf.float32)
4
5with tf.GradientTape() as tape:
6    y = tf.stack([
7        x[0] ** 2 + 3.0 * x[1],
8        x[0] * x[1]
9    ])
10
11jacobian = tape.jacobian(y, x)
12print(jacobian.numpy())

The expected Jacobian is:

[[2.0, 3.0],
 [2.0, 1.0]]

The first row comes from derivatives of x[0] ** 2 + 3 * x[1]. The second row comes from derivatives of x[0] * x[1].

Reading Tensor Shapes Correctly

Shape confusion is the most common problem. If the target has shape (2,) and the source has shape (2,), the Jacobian has shape (2, 2). If the target were shape (4, 3) and the source were shape (5,), the result would have shape (4, 3, 5).

This rule is consistent with the TensorFlow API: target dimensions first, source dimensions last.

jacobian Versus batch_jacobian

Use jacobian when you want the full derivative of the entire target with respect to the source.

Use batch_jacobian when each batch element should only depend on the corresponding input element. That avoids many zero entries and is often more efficient for batched models.

1import tensorflow as tf
2
3x = tf.constant([[1.0, 2.0], [3.0, 4.0]])
4
5with tf.GradientTape() as tape:
6    tape.watch(x)
7    y = x * x
8
9result = tape.batch_jacobian(y, x)
10print(result.shape)

For repeated Jacobian calculations, TensorFlow documentation notes that wrapping the code in @tf.function can improve performance because Jacobian computation uses vectorization machinery internally.

Watching Non-Variable Tensors

If the source is a tf.Variable, GradientTape watches it automatically. If it is an ordinary tensor, call tape.watch(source) yourself.

1x = tf.constant([1.0, 2.0])
2with tf.GradientTape() as tape:
3    tape.watch(x)
4    y = x * x
5print(tape.jacobian(y, x).numpy())

Without watch, the result may be None because TensorFlow did not record x as a differentiation source.

Reusing the Same Tape

If you need several Jacobians or a Jacobian plus a gradient from the same recorded computation, create the tape with persistent=True. Otherwise the tape is consumed after the first derivative call.

with tf.GradientTape(persistent=True) as tape:
    y = model(x)

Delete or let the tape go out of scope when you are done so the recorded graph can be released.

Sanity-Check the Result

When you are debugging a model, validate at least one Jacobian entry by hand or with a small finite-difference check. That catches indexing mistakes and shape misunderstandings early, especially when the target tensor is higher-rank.

Common Pitfalls

The first pitfall is confusing gradients with Jacobians. If the output is not scalar, gradient is not the same tool.

Another common problem is trying to reuse a non-persistent tape for multiple derivative calls. If you need more than one Jacobian from the same recorded computation, use persistent=True and release it afterward.

A third pitfall is building Jacobians for very large outputs without considering memory cost. The tensor can grow quickly.

Summary

• Use tf.GradientTape().jacobian(target, source) for vector-valued derivatives.
• The Jacobian shape is target.shape + source.shape.
• Use batch_jacobian for per-example derivatives in batched data.
• Call tape.watch(...) for non-Variable tensors.
• Watch memory usage when outputs or inputs are high-dimensional.