autodiff_impl.py

# Ref. http://www.cs.cmu.edu/~wcohen/10-605/notes/autodiff.pdf

from collections import namedtuple

Assignment = namedtuple("Item", ("z", "g", "y_list"))

# f(x_1, x_2) = (2x_1 + x_2)^2
l = [Assignment(z="z1", g="add", y_list=["x1", "x1"]),
     Assignment(z="z2", g="add", y_list=["z1", "x2"]),
     Assignment(z="f", g="square", y_list=["z2"])]

G = {"add": lambda a: a[0] + a[1],
     "square": lambda a: a[0] * a[0]}

# forward
val_input = {"x1": 3, "x2": 7}


def eval_forward(val, l, G):
    for assign in l:
        z, g, y = assign.z, assign.g, assign.y_list

        op = G[g]
        val[z] = op([val[y_i] for y_i in y])
    return val


print(f"Forward: {val_input} => {eval_forward(val_input, l, G)}")


# backpropagation
DG = {"add": [lambda a: 1, lambda a: 1],
      "square": [lambda a: 2 * a[0]]}


def backprop(d_value, val, DG, l):
    delta = {"f": d_value}
    for assign in reversed(l):
        z, g, y, k = assign.z, assign.g, assign.y_list, len(assign.y_list)

        for i in range(k):
            op = DG[g][i]
            if y[i] not in delta:
                delta[y[i]] = 0
            delta[y[i]] = delta[y[i]] + delta[z] * op([val[y_i] for y_i in y])
    return delta


d = 1
output = backprop(d, eval_forward(val_input, l, G), DG, l)
print(f"Back: f={d} => {output}")

output.txt

Forward: {'x1': 3, 'x2': 7} => {'x1': 3, 'x2': 7, 'z1': 6, 'z2': 13, 'f': 169}
Backprop: f=1 => {'f': 1, 'z2': 26, 'z1': 26, 'x2': 26, 'x1': 52}