ai-lab/m01_autograd/engine.py

"""AI01 lab: a scalar autograd engine. You write every method in this file.

Each docstring is a contract, and the tests check exactly these contracts.
Work through the lab steps in the book (AI01, section 3) in order, and run
`python check.py N` after finishing step N.

Ground rules
- Standard library only. `math` is all you need.
- Do not define __eq__ on Value. backward() keeps Values in a set, and a class
  that defines __eq__ without __hash__ becomes unhashable (bug B10 in AI01).
- Write each local derivative on paper before you type it (the Modular
  Workbench): what is d(out)/d(input), in terms of numbers you already have?
"""
import math


class Value:
    """One number in a computation, plus how much the final output depends on it.

    Attributes (the tests and tools/viz.py rely on these exact names):
      data       float. The number itself, computed in the forward pass.
      grad       float. d(final output)/d(this value). Starts at 0.0;
                 backward() fills it in.
      _prev      the Values this one was computed from (its children in the
                 graph), empty for a leaf. Store a tuple: it keeps the inputs
                 in order, so every run adds gradients in the same order and
                 prints the same digits. (A set also works, but it orders
                 Values by memory address, so the last digits drift between runs.)
      _op        str naming the operation that made this Value ('+', '*',
                 '**2', 'tanh', ...); '' for a leaf.
      label      str, an optional name used in pictures ('a', 'loss', ...).
      _backward  a function taking no arguments. Calling it adds this Value's
                 contribution to each child's grad (local derivative times
                 self.grad). For a leaf it does nothing.
    """

    # ------------------------------------------------------------------
    # STEP 1: hold a number and build the graph going forward.
    # ------------------------------------------------------------------
    def __init__(self, data, _children=(), _op="", label=""):
        # TODO(human): Step 1. Store data (as a float) and start grad at 0.0.
        # Record _prev, _op and label, and give the node a _backward that does nothing.
        raise NotImplementedError("Step 1: Value.__init__")

    def __repr__(self):
        """Return text like 'Value(data=2.5, grad=0.0)'."""
        raise NotImplementedError("Step 1: Value.__repr__")

    def __add__(self, other):
        """self + other.

        If `other` is a plain int or float, wrap it in a Value first.
        Forward (step 1): return a new Value with data self.data + other.data,
        children (self, other), and op '+'.
        Backward (step 2): attach a _backward to the result. The local
        derivative of a + b with respect to each input is 1.
        """
        raise NotImplementedError("Step 1: Value.__add__")

    def __mul__(self, other):
        """self * other. Wrap plain numbers; op '*'.

        Backward (step 2): the local derivative of a * b with respect to a is
        b's data, and with respect to b it is a's data.
        """
        raise NotImplementedError("Step 1: Value.__mul__")

    # ------------------------------------------------------------------
    # STEP 3: the whole backward pass.
    # ------------------------------------------------------------------
    def backward(self):
        """Fill in .grad for every Value this one depends on.

        1. Put every node of the graph in topological order: each node after
           all of its children. Use a depth-first search with a visited set.
        2. Seed: the output's gradient with respect to itself is 1.
        3. Visit the nodes in REVERSE topological order and call each _backward.

        Assumes every grad in the graph starts at 0.0. To backpropagate again,
        zero the grads first (or build a fresh graph).

        A recursive search is fine for steps 1-8. The optional stretch test
        (python check.py --stretch) uses a 5000-node chain, deeper than
        Python's limit of about 1000 nested calls, so it needs an explicit stack.
        """
        raise NotImplementedError("Step 3: Value.backward")

    # ------------------------------------------------------------------
    # STEP 5: more operations. Each needs a forward value AND a _backward.
    # Reuse the operations you already have wherever you can.
    # ------------------------------------------------------------------
    def __pow__(self, n):
        """self ** n, where n is a plain int or float (not a Value). Op f'**{n}'.

        Local derivative: n * self.data ** (n - 1). Watch the sign of n - 1
        when n is negative (a classic slip: -48x^4 written for -48x^(-4)).
        """
        raise NotImplementedError("Step 5: Value.__pow__")

    def __neg__(self):
        """-self."""
        raise NotImplementedError("Step 5: Value.__neg__")

    def __radd__(self, other):
        """other + self, when `other` is a plain number (for example 2 + a)."""
        raise NotImplementedError("Step 5: Value.__radd__")

    def __sub__(self, other):
        """self - other."""
        raise NotImplementedError("Step 5: Value.__sub__")

    def __rsub__(self, other):
        """other - self, when `other` is a plain number (for example 5 - a)."""
        raise NotImplementedError("Step 5: Value.__rsub__")

    def __rmul__(self, other):
        """other * self, when `other` is a plain number (for example 2 * a)."""
        raise NotImplementedError("Step 5: Value.__rmul__")

    def __truediv__(self, other):
        """self / other. Hint: a / b = a * b**(-1)."""
        raise NotImplementedError("Step 5: Value.__truediv__")

    def __rtruediv__(self, other):
        """other / self, when `other` is a plain number (for example 4 / a)."""
        raise NotImplementedError("Step 5: Value.__rtruediv__")

    def exp(self):
        """e ** self. Op 'exp'. Local derivative: e ** x, which is out.data itself."""
        raise NotImplementedError("Step 5: Value.exp")

    def tanh(self):
        """Hyperbolic tangent. Op 'tanh'. Local derivative: 1 - tanh(x)**2.

        Must work for large inputs: tanh of 1000 is 1.0 with gradient 0.0,
        and tanh of -1000 is -1.0, with no OverflowError.
        """
        raise NotImplementedError("Step 5: Value.tanh")

    def relu(self):
        """max(0, x). Op 'relu'.

        Local derivative: 1 where x > 0, and 0 where x <= 0. At exactly 0 the
        true derivative does not exist; 0 is a convention (PyTorch uses it too).
        """
        raise NotImplementedError("Step 5: Value.relu")