ai-lab/m01_autograd/tests/test_gradcheck.py

"""AI01 lab tests for gradcheck.py (step 6)."""
import unittest

from _impl import load

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


class Step6GradCheck(unittest.TestCase):
    def test_partials_of_the_worked_example(self):
        # f(x, y) = x^2 y + 3y at (1, 3): df/dx = 2xy = 6, df/dy = x^2 + 3 = 4.
        f = lambda v: v[0] ** 2 * v[1] + 3 * v[1]
        g = gradcheck.numerical_grad(f, [1.0, 3.0])
        self.assertEqual(len(g), 2)
        self.assertAlmostEqual(g[0], 6.0, places=6)
        self.assertAlmostEqual(g[1], 4.0, places=6)

    def test_uses_the_central_difference(self):
        # For x^3 at x = 1 with h = 1e-3, the central difference is off by
        # h^2 = 1e-6, but the forward difference is off by about 3h = 3e-3.
        g = gradcheck.numerical_grad(lambda v: v[0] ** 3, [1.0], h=1e-3)
        self.assertLess(abs(g[0] - 3.0), 1e-5)

    def test_uses_the_h_it_is_given(self):
        # The central difference of x^3 at 1 is exactly 3 + h^2, so h = 0.1 gives 3.01.
        g = gradcheck.numerical_grad(lambda v: v[0] ** 3, [1.0], h=0.1)
        self.assertAlmostEqual(g[0], 3.01, places=9)

    def test_leaves_the_point_unchanged(self):
        xs = [1.0, 3.0]
        gradcheck.numerical_grad(lambda v: v[0] * v[1], xs)
        self.assertEqual(xs, [1.0, 3.0])

    def test_other_coordinates_are_held_fixed(self):
        # f = x*y*z at (2, 3, 4): partials are yz = 12, xz = 8, xy = 6.
        g = gradcheck.numerical_grad(lambda v: v[0] * v[1] * v[2], [2.0, 3.0, 4.0])
        for got, want in zip(g, [12.0, 8.0, 6.0]):
            self.assertAlmostEqual(got, want, places=5)

    def test_max_grad_error_is_tiny_for_a_correct_engine(self):
        f = lambda v: (v[0] * v[1] + v[1]) * v[2]      # the running example
        err = gradcheck.max_grad_error(f, [3.0, -2.0, 5.0])
        self.assertLess(err, 1e-5)

    def test_max_grad_error_catches_a_broken_gradient(self):
        def broken(v):
            out = v[0] * 1.0
            out._backward = lambda: None     # gradient never reaches v[0]
            return out + v[1]
        err = gradcheck.max_grad_error(broken, [2.0, 5.0])
        self.assertAlmostEqual(err, 1.0, places=5)

    def test_max_grad_error_catches_a_gradient_that_is_too_big(self):
        def doubled(v):
            out = v[0] * 1.0
            def _backward():
                v[0].grad += 2.0 * out.grad  # twice the true gradient
            out._backward = _backward
            return out + v[1]
        err = gradcheck.max_grad_error(doubled, [2.0, 5.0])
        self.assertAlmostEqual(err, 1.0, places=5)

    def test_max_grad_error_passes_h_on(self):
        # Backprop gives exactly 3 for x^3 at 1; the central difference with h = 0.1 gives 3.01.
        err = gradcheck.max_grad_error(lambda v: v[0] ** 3, [1.0], h=0.1)
        self.assertAlmostEqual(err, 0.01, places=9)


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