ai-lab/m01_autograd/tests/test_engine.py

"""AI01 lab tests for engine.py (steps 1-6).

Run from m01_autograd/ with `python check.py N`, which runs steps 1..N.
Each test's name says what it checks; the book's step N explains why.
"""
import math
import os
import random
import unittest

from _impl import load

engine = load("engine")
Value = engine.Value


def running_example():
    """The chapter's running example: L = (a*b + b) * c at a=3, b=-2, c=5."""
    a, b, c = Value(3.0, label="a"), Value(-2.0, label="b"), Value(5.0, label="c")
    e = a * b          # -6
    d = e + b          # -8
    L = d * c          # -40
    return a, b, c, e, d, L


def contains(nodes, target):
    """True if `target` is one of `nodes`, by identity (never by ==)."""
    return any(node is target for node in nodes)


# ----------------------------------------------------------------------
class Step1Forward(unittest.TestCase):
    def test_holds_data_and_starts_with_zero_grad(self):
        v = Value(2.5)
        self.assertEqual(v.data, 2.5)
        self.assertEqual(v.grad, 0.0)
        self.assertEqual(len(v._prev), 0)

    def test_data_is_stored_as_float(self):
        self.assertIsInstance(Value(3).data, float)

    def test_label_is_stored(self):
        self.assertEqual(Value(1.0, label="a").label, "a")

    def test_add_and_mul_return_new_values(self):
        a, b = Value(3.0), Value(-2.0)
        s, p = a + b, a * b
        self.assertIsInstance(s, Value)
        self.assertIsInstance(p, Value)
        self.assertEqual(s.data, 1.0)
        self.assertEqual(p.data, -6.0)
        self.assertEqual((a.data, b.data), (3.0, -2.0), "inputs must not change")

    def test_plain_numbers_on_the_right_are_wrapped(self):
        a = Value(3.0)
        self.assertEqual((a + 1).data, 4.0)
        self.assertEqual((a * 2).data, 6.0)

    def test_children_and_op_are_recorded(self):
        a, b = Value(3.0), Value(-2.0)
        s, p = a + b, a * b
        self.assertTrue(contains(s._prev, a) and contains(s._prev, b))
        self.assertTrue(contains(p._prev, a) and contains(p._prev, b))
        self.assertEqual(s._op, "+")
        self.assertEqual(p._op, "*")
        self.assertEqual(a._op, "")

    def test_running_example_forward(self):
        *_, e, d, L = running_example()
        self.assertEqual((e.data, d.data, L.data), (-6.0, -8.0, -40.0))

    def test_repr_shows_data_and_grad(self):
        text = repr(Value(2.5))
        self.assertIn("2.5", text)
        self.assertIn("grad", text)

    def test_values_are_hashable(self):
        # backward() keeps Values in a set. Defining __eq__ breaks this (bug B10).
        a = Value(1.0)
        self.assertIn(a, {a})


# ----------------------------------------------------------------------
class Step2LocalBackward(unittest.TestCase):
    """Each node's _backward, called by hand. No backward() yet."""

    def test_add_passes_the_gradient_through_unchanged(self):
        a, b = Value(3.0), Value(-2.0)
        s = a + b
        s.grad = 4.0
        s._backward()
        self.assertEqual((a.grad, b.grad), (4.0, 4.0))

    def test_mul_sends_each_input_the_other_inputs_value(self):
        a, b = Value(3.0), Value(-2.0)
        p = a * b
        p.grad = 4.0
        p._backward()
        self.assertEqual(a.grad, -8.0)   # b.data * p.grad
        self.assertEqual(b.grad, 12.0)   # a.data * p.grad

    def test_leaf_backward_does_nothing(self):
        a = Value(1.0)
        a.grad = 7.0
        a._backward()
        self.assertEqual(a.grad, 7.0)

    def test_wrapped_constant_gets_the_chain_rule_too(self):
        a = Value(3.0)
        p = a * 2
        p.grad = 1.5
        p._backward()
        self.assertEqual(a.grad, 3.0)

    def test_manual_backprop_through_the_running_example(self):
        a, b, c, e, d, L = running_example()
        L.grad = 1.0
        L._backward()          # L = d * c
        d._backward()          # d = e + b
        e._backward()          # e = a * b
        self.assertEqual((d.grad, c.grad), (5.0, -8.0))
        self.assertEqual(e.grad, 5.0)
        self.assertEqual(a.grad, -10.0)
        self.assertEqual(b.grad, 20.0, "b is used twice: 5 from d plus 15 from e")


# ----------------------------------------------------------------------
class Step3Backward(unittest.TestCase):
    def test_running_example_gradients(self):
        a, b, c, e, d, L = running_example()
        L.backward()
        self.assertEqual(L.grad, 1.0)
        self.assertEqual((d.grad, e.grad), (5.0, 5.0))
        self.assertEqual((a.grad, b.grad, c.grad), (-10.0, 20.0, -8.0))

    def test_output_is_seeded_with_one(self):
        x = Value(2.0)
        y = x * 3
        y.backward()
        self.assertEqual(y.grad, 1.0)
        self.assertEqual(x.grad, 3.0)

    def test_a_lone_leaf(self):
        x = Value(4.0)
        x.backward()
        self.assertEqual(x.grad, 1.0)

    def test_long_chain_is_processed_in_reverse_topological_order(self):
        x = Value(1.0)
        y = x
        for _ in range(30):
            y = y * 2 + 1
        y.backward()
        self.assertEqual(x.grad, 2.0 ** 30)

    def test_shared_interior_node_is_processed_once(self):
        a = Value(2.0)
        t = a * 3                 # an interior node...
        y = t * t + t             # ...used three times
        y.backward()
        self.assertEqual(t.grad, 13.0)   # dy/dt = 2t + 1 at t = 6
        self.assertEqual(a.grad, 39.0)   # dy/da = 13 * 3

    def test_diamond(self):
        x = Value(1.5)
        left, right = x * 2, x + 4
        y = left * right          # (2x)(x + 4) = 2x^2 + 8x
        y.backward()
        self.assertEqual(x.grad, 4 * 1.5 + 8)


# ----------------------------------------------------------------------
class Step4Accumulate(unittest.TestCase):
    def test_a_plus_a(self):
        a = Value(3.0)
        (a + a).backward()
        self.assertEqual(a.grad, 2.0)

    def test_a_times_a(self):
        a = Value(3.0)
        (a * a).backward()
        self.assertEqual(a.grad, 6.0)

    def test_cube_by_repeated_multiplication(self):
        x = Value(3.0)
        (x * x * x).backward()
        self.assertEqual(x.grad, 27.0)

    def test_one_input_many_uses(self):
        x = Value(2.0)
        total = x * 0
        for _ in range(10):
            total = total + x
        total.backward()
        self.assertEqual(x.grad, 10.0)


# ----------------------------------------------------------------------
class Step5MoreOps(unittest.TestCase):
    """Each op is tested inside a chain (times 3), so a _backward that forgets
    to multiply by out.grad cannot pass (bug B7)."""

    def check(self, f, x0, value, slope, places=10):
        x = Value(x0)
        y = f(x) * 3
        y.backward()
        self.assertAlmostEqual(y.data, 3 * value, places=places)
        self.assertAlmostEqual(x.grad, 3 * slope, places=places)

    def test_pow_positive_integer(self):
        self.check(lambda x: x ** 3, 2.0, 8.0, 12.0)

    def test_pow_negative_exponent(self):
        # d/dx x^(-2) = -2 x^(-3). The exponent stays negative.
        self.check(lambda x: x ** -2, 2.0, 0.25, -0.25)

    def test_pow_fractional_exponent(self):
        self.check(lambda x: x ** 0.5, 4.0, 2.0, 0.25)

    def test_neg(self):
        self.check(lambda x: -x, 3.0, -3.0, -1.0)

    def test_sub(self):
        self.check(lambda x: x - 5, 3.0, -2.0, 1.0)

    def test_rsub(self):
        self.check(lambda x: 5 - x, 3.0, 2.0, -1.0)

    def test_radd(self):
        self.check(lambda x: 2 + x, 3.0, 5.0, 1.0)

    def test_rmul(self):
        self.check(lambda x: 2 * x, 3.0, 6.0, 2.0)

    def test_div_by_number(self):
        self.check(lambda x: x / 4, 2.0, 0.5, 0.25)

    def test_rdiv(self):
        # 4/x = 4x^(-1), so the slope is -4x^(-2).
        self.check(lambda x: 4 / x, 2.0, 2.0, -1.0)

    def test_value_divided_by_value(self):
        a, b = Value(3.0), Value(2.0)
        q = a / b
        q.backward()
        self.assertAlmostEqual(q.data, 1.5)
        self.assertAlmostEqual(a.grad, 0.5)     # 1/b
        self.assertAlmostEqual(b.grad, -0.75)   # -a/b^2

    def test_sub_of_two_values(self):
        a, b = Value(3.0), Value(2.0)
        (a - b).backward()
        self.assertEqual((a.grad, b.grad), (1.0, -1.0))

    def test_exp(self):
        self.check(lambda x: x.exp(), 1.0, math.e, math.e)

    def test_tanh(self):
        t = math.tanh(0.5)
        self.check(lambda x: x.tanh(), 0.5, t, 1 - t * t)

    def test_tanh_saturates_without_overflow(self):
        for x0, expected in ((1000.0, 1.0), (-1000.0, -1.0)):
            x = Value(x0)
            y = x.tanh()
            y.backward()
            self.assertEqual(y.data, expected)
            self.assertEqual(x.grad, 0.0)

    def test_relu(self):
        self.check(lambda x: x.relu(), 3.0, 3.0, 1.0)
        self.check(lambda x: x.relu(), -2.0, 0.0, 0.0)

    def test_relu_gradient_at_zero_is_zero(self):
        self.check(lambda x: x.relu(), 0.0, 0.0, 0.0)

    def test_new_ops_record_children(self):
        x = Value(2.0)
        for out in (x ** 2, x.exp(), x.tanh(), x.relu()):
            self.assertTrue(contains(out._prev, x))
            self.assertIsInstance(out._op, str)
            self.assertNotEqual(out._op, "")


# ----------------------------------------------------------------------
def _random_plan(rng, n_inputs, n_ops):
    """A random computation that reuses earlier results, so fan-out is common."""
    kinds = ["+", "*", "-", "/", "c+", "c*", "c-", "c/",
             "tanh", "exp", "sq", "cube", "inv", "neg"]
    plan = []
    for k in range(n_ops):
        avail = n_inputs + k
        plan.append((rng.choice(kinds), rng.randrange(avail), rng.randrange(avail),
                     rng.choice([-2.0, -0.5, 0.5, 1.5, 3.0])))
    return plan


def _run_plan(inputs, plan):
    nodes = list(inputs)
    for kind, i, j, c in plan:
        a, b = nodes[i], nodes[j]
        if kind == "+":
            r = a + b
        elif kind == "*":
            r = a * b
        elif kind == "-":
            r = a - b
        elif kind == "/":
            r = a / (b * b + 1)          # never divides by zero
        elif kind == "c+":
            r = c + a
        elif kind == "c*":
            r = c * a
        elif kind == "c-":
            r = c - a
        elif kind == "c/":
            r = c / (a * a + 1)
        elif kind == "tanh":
            r = a.tanh()
        elif kind == "exp":
            r = a.tanh().exp()           # bounded, so nothing overflows
        elif kind == "sq":
            r = a ** 2
        elif kind == "cube":
            r = a.tanh() ** 3
        elif kind == "inv":
            r = (a * a + 1) ** -1
        else:
            r = -a
        nodes.append(r)
    return nodes[-1] + nodes[-2] * nodes[-3]


class Step6RandomGradcheck(unittest.TestCase):
    """Backprop must agree with finite differences on 40 random graphs.
    This test does its own finite differences; it does not use your gradcheck.py."""

    def test_random_graphs_match_central_differences(self):
        rng = random.Random(20260924)
        h = 1e-6
        for trial in range(40):
            n_inputs = rng.randint(2, 4)
            plan = _random_plan(rng, n_inputs, rng.randint(4, 10))
            xs = [rng.uniform(-1.5, 1.5) for _ in range(n_inputs)]

            leaves = [Value(x) for x in xs]
            out = _run_plan(leaves, plan)
            out.backward()

            for i in range(n_inputs):
                up, down = list(xs), list(xs)
                up[i] += h
                down[i] -= h
                f_up = _run_plan([Value(v) for v in up], plan).data
                f_down = _run_plan([Value(v) for v in down], plan).data
                numeric = (f_up - f_down) / (2 * h)
                with self.subTest(trial=trial, input=i, plan=plan):
                    self.assertLessEqual(abs(leaves[i].grad - numeric),
                                         1e-5 * max(1.0, abs(numeric)))


# ----------------------------------------------------------------------
@unittest.skipUnless(os.environ.get("AILAB_STRETCH"), "stretch test: python check.py --stretch")
class StretchDeepGraph(unittest.TestCase):
    def test_backward_on_a_5000_node_chain(self):
        # Python stops recursing at about 1000 nested calls (sys.getrecursionlimit()).
        # A recursive topological sort fails here; an explicit stack does not.
        x = Value(1.0)
        y = x
        for _ in range(5000):
            y = y + 1
        y.backward()
        self.assertEqual(x.grad, 1.0)


if __name__ == "__main__":
    unittest.main()