Autograd Basics - PyTorch Basics

Autograd Basics

What you'll build: two gradient computations: the derivative of a scalar polynomial, and the backward pass of a linear layer Y = X @ W + b.

Background

Autograd records every operation of the forward pass in a graph, then walks it backwards applying the chain rule. In PyTorch you mark inputs with requires_grad=True, call .backward() and read .grad:

x = torch.tensor(2.0, requires_grad=True)
y = x ** 2 + 3 * x + 1
y.backward()
x.grad  # tensor(7.) because dy/dx = 2x + 3

The PyTorch track uses autograd exactly like this. The NumPy track writes the same gradients by hand, which is what backward() computes under the hood and a favourite interview question.

Formulas

1. Scalar function. For y = x² + 3x + 1: dy/dx = 2x + 3.

2. Linear layer. Forward: Y = X @ W + b with

Symbol	Meaning	Shape
X	input batch	(batch, in_features)
W	weights	(in_features, out_features)
b	bias, broadcast over the batch	(out_features,)
G = grad_output	upstream gradient dL/dY	(batch, out_features)

Backward, given G:

dL/dW = X.T @ G          shape (in_features, out_features)
dL/db = G.sum(axis=0)    shape (out_features,)   (b was broadcast over the batch)
dL/dX = G @ W.T          shape (batch, in_features)

A quick check that never fails: every gradient has the same shape as the tensor it belongs to, and there is only one way to multiply the available matrices to get that shape.

Expected Return Format

• compute_gradients(x) returns {'y': float, 'grad': float}: the forward value and dy/dx.
• linear_layer_gradients(X, W, b, grad_output) returns a dict with 'grad_W', 'grad_b' and 'grad_X', shaped like W, b and X.

Constraints

• NumPy track: compute the gradients analytically with the formulas above.
• PyTorch track: let autograd do it (requires_grad=True, then .backward(); for the non-scalar Y, call Y.backward(grad_output)).

Follow-up

• Why does Y.backward() without an argument fail when Y is not a scalar?
• What is dL/dX used for, if X is the input data and is never updated?

Input: compute_gradients(2.0)
Output: {'y': 11.0, 'grad': 7.0}

Input: X = [[1, 2], [3, 4]], W = [[1, 0, -1], [2, 1, 0]], b = [0, 0, 0], grad_output = ones((2, 3))
Output: {'grad_W': [[4, 4, 4], [6, 6, 6]], 'grad_b': [2, 2, 2], 'grad_X': [[0, 3], [0, 3]]}

import numpy as np

def compute_gradients(x: float) -> dict:
    """
    Compute the forward value and derivative of y = x^2 + 3x + 1.

    Args:
        x: Input value

    Returns:
        dict with 'y' (forward result) and 'grad' (dy/dx), both floats
    """
    # Your code here
    pass


def linear_layer_gradients(X: np.ndarray, W: np.ndarray, b: np.ndarray,
                           grad_output: np.ndarray) -> dict:
    """
    Compute the gradients of a linear layer Y = X @ W + b.

    Args:
        X: Input (batch, in_features)
        W: Weights (in_features, out_features)
        b: Bias (out_features,)
        grad_output: Upstream gradient dL/dY (batch, out_features)

    Returns:
        dict with 'grad_W' (in_features, out_features), 'grad_b'
        (out_features,) and 'grad_X' (batch, in_features)
    """
    # Your code here
    pass