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:

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:

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

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?