4. Tensor representation and automatic differentiation

Every TensorFlow program ultimately manipulates two kinds of values: immutable tensors that carry data through the graph, and mutable variables that hold the parameters a model learns. What makes TensorFlow more than a numerical library, however, is its ability to record the operations applied to those values and replay them backward to compute derivatives. This chapter examines the tensor data model, the role of tf.Variable, the mechanics of tf.GradientTape, and how to override the default differentiation behavior with custom gradients.

4.1 Tensors as Immutable, Typed, Shaped Arrays

A tf.Tensor is an n-dimensional array with a fixed dtype and a shape. Unlike a NumPy array, a tensor is immutable: you never assign into a tensor's elements. Instead, operations produce new tensors. This immutability is what allows TensorFlow to build a symbolic record of computation and to optimize or differentiate it.

Three attributes matter most in practice:

• Rank — the number of dimensions. A scalar has rank 0, a vector rank 1, a matrix rank 2.
• Shape — the size along each dimension. A dimension may be None in a tf.function signature, meaning "unknown at trace time."
• dtype — for example tf.float32, tf.int64, or tf.bool. Gradient computation requires floating-point types; integer tensors are treated as constants.

import tensorflow as tf

x = tf.constant([[1.0, 2.0], [3.0, 4.0]], dtype=tf.float32)
print(x.shape)   # (2, 2)
print(x.dtype)   # <dtype: 'float32'>
print(tf.rank(x))  # tf.Tensor(2, shape=(), dtype=int32)

# Immutability: this creates a new tensor, it does not modify x
y = x + 10.0

When you need to change values in place — for example, to update weights during training — you use a tf.Variable instead.

4.2 tf.Variable: Mutable State with a Traceable Identity

A tf.Variable wraps a tensor and exposes mutation methods such as assign, assign_add, and assign_sub. Variables are the natural home for model parameters, optimizer accumulators, and counters. Two properties distinguish them from plain tensors:

1. They persist across calls to a tf.function, so their values survive between training steps.
2. They are tracked by the tape when read inside a GradientTape context, provided the tape is watching them.

w = tf.Variable(tf.random.normal([3, 2]), name="weights")
b = tf.Variable(tf.zeros([2]), name="bias")

w.assign_add(tf.ones_like(w))   # in-place update
print(w.trainable)              # True by default

By default a variable is trainable, meaning GradientTape will watch it automatically. Setting trainable=False is useful for frozen layers or for statistics such as batch-norm moving averages that should not receive gradients.

4.3 GradientTape: Recording and Replaying Computation

tf.GradientTape is a context manager that records every differentiable operation applied to watched tensors. After the forward computation completes, you call tape.gradient(target, sources) to obtain the derivative of a scalar target with respect to one or more sources. The tape is consumed by that call; a second call raises an error unless you requested a persistent tape.

x = tf.Variable(3.0)
with tf.GradientTape() as tape:
    y = x * x + 2 * x + 1     # y = x^2 + 2x + 1
dy_dx = tape.gradient(y, x)   # 2x + 2 = 8 at x = 3
print(dy_dx.numpy())          # 8.0

Three rules govern correct usage:

• Watch before you compute. Constants are not watched by default. Call tape.watch(t) on a constant tensor before using it in the recorded block.
• Target must be scalar (or you must supply a gradient). For a non-scalar target, pass output_gradients to specify the upstream vector.
• One tape, one gradient call unless persistent=True is set, in which case you must delete the tape explicitly.

For higher-order derivatives, nest tapes. The inner tape must be created inside the outer tape's context so that the gradient computation itself is recorded.

x = tf.Variable(2.0)
with tf.GradientTape() as outer:
    with tf.GradientTape() as inner:
        y = x ** 3
    dy_dx = inner.gradient(y, x)      # 3x^2
d2y_dx2 = outer.gradient(dy_dx, x)    # 6x = 12

4.4 Custom Gradients with @tf.custom_gradient

Sometimes the default derivative is numerically unstable, undefined at a point, or simply not what you want — for example, a straight-through estimator for quantization, or a numerically stable log-sum-exp. The @tf.custom_gradient decorator lets you define both the forward function and a gradient function that receives the upstream gradient and returns gradients with respect to each input.

@tf.custom_gradient
def log1p_exp(x):
    # forward: log(1 + exp(x)), stable for large x
    y = tf.math.log1p(tf.exp(-tf.abs(x))) + tf.maximum(x, 0.0)

    def grad(upstream):
        # d/dx log(1+exp(x)) = sigmoid(x)
        sig = tf.sigmoid(x)
        return upstream * sig

    return y, grad

The gradient function must return one tensor per input of the forward function, in the same order. If the forward function has non-tensor arguments, they are passed through unchanged. A common mistake is to return a list whose length does not match the number of inputs, which produces a confusing shape error at gradient time.

4.5 A Minimal Training Step

The following example ties the pieces together: a variable parameter, a tape recording a loss, and an optimizer applying the gradient. This is the canonical pattern used inside tf.function for performance.

w = tf.Variable(1.0)
optimizer = tf.keras.optimizers.SGD(learning_rate=0.1)

for step in range(5):
    with tf.GradientTape() as tape:
        loss = (w - 3.0) ** 2      # minimum at w = 3
    grad = tape.gradient(loss, w)
    optimizer.apply_gradients([(grad, w)])
    print(step, float(w), float(loss))

Expected behavior: w moves from 1.0 toward 3.0, and the loss decreases monotonically. If w does not change, the usual causes are (a) the variable was created with trainable=False, (b) the tape was created before the variable was read, or (c) the gradient is None because the loss does not depend on w through differentiable operations.

4.6 How the Pieces Fit Together

The diagram below summarizes the data and gradient flow. Tensors flow forward through operations; the tape records each operation; the backward pass traverses the recorded graph in reverse, multiplying local Jacobians to produce gradients that the optimizer uses to update variables.

The key insight is that the tape does not store values — it stores the graph of operations and the intermediate tensors needed to compute local derivatives. This is why memory grows with the number of recorded operations, and why releasing the tape promptly (by exiting the context) matters for long training loops.

4.7 Summary

Tensors are immutable, typed, shaped arrays; variables add mutable, tracked state. GradientTape records differentiable operations and replays them in reverse to produce gradients, and @tf.custom_gradient lets you replace the default derivative when numerical or semantic needs demand it. Mastering these four constructs — tensor, variable, tape, custom gradient — is the prerequisite for every training loop and every custom layer you will write in the rest of this course.