Implement a minimal reverse-mode autograd engine on PyTorch tensors

In an ML coding interview you are asked to implement the core of an automatic differentiation ("autograd") system of the kind PyTorch provides, with the PyTorch library available in the environment. Build a small `Tensor` class whose operations record how each result was computed, and a `backward()` method that computes, by reverse-mode differentiation, the gradient of a scalar output with respect to every tensor that requires gradients. The exact interface and set of operations were not specified, so assume this minimal version and confirm it with the interviewer: ```python import torch class Tensor: def __init__(self, data: torch.Tensor, requires_grad: bool = False): ... # self.data: torch.Tensor the stored values # self.grad: torch.Tensor | None filled in by backward() def __add__(self, other: "Tensor") -> "Tensor": ... # elementwise def __mul__(self, other: "Tensor") -> "Tensor": ... # elementwise def __matmul__(self, other: "Tensor") -> "Tensor": ... # 2-D matrix product def relu(self) -> "Tensor": ... def sum(self) -> "Tensor": ... # sum of all elements def backward(self) -> None: ... # called on a one-element output ``` PyTorch tensors serve only as array storage: the gradients must come from your own backward rules, not from PyTorch's built-in autograd. You may use PyTorch's autograd in tests to check your results. For example: ```python x = Tensor(torch.tensor([[1.0, 2.0]]), requires_grad=True) # shape (1, 2) W = Tensor(torch.tensor([[1.0, -1.0], [0.5, -2.0]]), requires_grad=True) # shape (2, 2) y = (x @ W).relu().sum() # x @ W = [[2.0, -5.0]], so y = 2.0 y.backward() # x.grad == [[1.0, 0.5]] # W.grad == [[1.0, 0.0], [2.0, 0.0]] a = Tensor(torch.tensor(3.0), requires_grad=True) b = a * a + a # b = 12.0 b.backward() # a.grad == 7.0, because db/da = 2a + 1 ``` ```hint What a result remembers When an operation creates a new tensor, decide what that tensor must keep so that, later, it can turn the gradient of its own value into gradients for its inputs. ``` ```hint Order of the backward pass A tensor can feed several later operations. Think about which order of visiting the graph guarantees that a tensor's gradient is complete before you pass it on to its inputs. ``` ### Constraints and Clarifications - The forward computation builds a directed acyclic graph. A tensor may be used by several operations, or twice by the same operation (as in `a * a`). - `@` takes two 2-D tensors of shapes `(n, m)` and `(m, p)`. - Every gradient has the same shape as the tensor it belongs to. - A result requires gradients when at least one of its inputs does. ### Clarifying Questions - Should the engine be a standalone graph over PyTorch tensors, as assumed above, or should each operation be written as a custom `torch.autograd.Function` with its own `forward` and `backward`? The task says only that the autograd machinery is written with the PyTorch library. - Must `+` and `*` broadcast, for example when adding a bias vector to every row of a batch? - Is `backward()` only ever called on a one-element output, or must it accept an upstream gradient for a non-scalar output? - If `backward()` runs twice, should leaf gradients accumulate as in PyTorch, or be overwritten? ### What a Strong Answer Covers - A graph representation in which each result records its inputs and a local backward rule - Correct local gradients for every operation, including the transposes in the matrix product and the ReLU mask, with each gradient matching its tensor's shape - A backward pass in a valid order that sums the contributions for tensors used more than once - Correct handling of tensors that do not require gradients, and no reliance on PyTorch's own autograd for the gradients - A test strategy that compares results with PyTorch's autograd or with finite differences ### Follow-up Questions - How would you support broadcasting in `+` and `*`, and what must the backward pass do to the gradient of a broadcast operand? - How would you implement a `no_grad` mode, and where does it save memory during inference? - Your graph has hundreds of thousands of nodes. What breaks in a recursive implementation, and how do you release memory that is no longer needed after `backward()`? - How would you add softmax with cross-entropy as a single operation, and why is that better than composing `exp`, `sum` and `log`?