Add backward-mode autodiff

.gitignore

@@ -158,3 +158,5 @@ cython_debug/
 #  and can be added to the global gitignore or merged into this file.  For a more nuclear
 #  option (not recommended) you can uncomment the following to ignore the entire idea folder.
 #.idea/
+
+uv.lock

main.py

@@ -0,0 +1,25 @@
+from simpleautodiff import *
+
+def main():
+    Node.verbose = True
+
+    # create root nodes
+    x1 = Node(2)
+    x2 = Node(5)
+
+    # create computational graph and evaluate function value
+    y = sub(add(log(x1), mul(x1, x2)), sin(x2))
+    # perform forward-mode autodiff
+    print("\n--- Forward mode (root = x1) ---")
+    forward(x1)
+    print("\n--- Forward mode (root = x2) ---")
+    forward(x2)
+
+    # perform reverse-mode autodiff
+    print("\n--- Reverse mode (output = y) ---")
+    backward(y)
+    print("\ndy/dx1 =", x1.partial_derivative.__round__(3))
+    print("dy/dx2 =", x2.partial_derivative.__round__(3))
+
+if __name__ == "__main__":
+    main()

pyproject.toml

@@ -0,0 +1,14 @@
+[project]
+name = "simpleautodiff"
+version = "0.1.0"
+description = "Simple autodiff demo"
+readme = "README.md"
+requires-python = ">=3.12"
+
+dependencies = [
+    "numpy>=1.26.0",
+    "scipy>=1.12.0",
+    "matplotlib>=3.8.0",
+    "pandas>=2.2.0",
+]
+

simpleautodiff/simpleautodiff.py

@@ -92,6 +92,17 @@ def add_children(node):
 
 
 def forward(rootNode):
+    # traverse the full graph in both directions to reset stale gradients
+    all_nodes, stack = set(), [rootNode]
+    while stack:
+        node = stack.pop()
+        if node not in all_nodes:
+            all_nodes.add(node)
+            stack.extend(node.parent_nodes)
+            stack.extend(node.child_nodes)
+    for node in all_nodes:
+        node.partial_derivative = 0
+
     rootNode.partial_derivative = 1
     ordering = topological_order(rootNode)
     for node in ordering[1:]:
@@ -119,3 +130,38 @@ def forward(rootNode):
                 value_process.strip(" + "),
                 str(node.partial_derivative.__round__(3)))
             )
+
+def reverse_topological_order(outputNode):
+    def add_parents(node):
+        if node not in visited:
+            visited.add(node)
+            for parent in node.parent_nodes:
+                add_parents(parent)
+            ordering.append(node)
+    ordering, visited = [], set()
+    add_parents(outputNode)
+    return list(reversed(ordering))
+
+
+def backward(outputNode):
+    ordering = reverse_topological_order(outputNode)
+    for node in ordering:
+        node.partial_derivative = 0
+    outputNode.partial_derivative = 1
+    for node in ordering:
+        for i, parent in enumerate(node.parent_nodes):
+            dnode_dparent = node.grad_wrt_parents[i]
+            parent.partial_derivative += dnode_dparent * node.partial_derivative
+
+        if Node.verbose:
+            for i, parent in enumerate(node.parent_nodes):
+                dnode_dparent = node.grad_wrt_parents[i]
+                print('d{:<2}/d{:<2} += (d{}/d{})*(d{}/d{}) = ({})({}) = {:<5}'.format(
+                    outputNode.name, parent.name,
+                    outputNode.name, node.name,
+                    node.name, parent.name,
+                    str(node.partial_derivative.__round__(3)),
+                    str(dnode_dparent.__round__(3)),
+                    str(parent.partial_derivative.__round__(3)))
+                )
+