diff --git a/.gitignore b/.gitignore index 68bc17f..9e019b8 100644 --- a/.gitignore +++ b/.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 diff --git a/main.py b/main.py new file mode 100644 index 0000000..fad6374 --- /dev/null +++ b/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() diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..c3f70d4 --- /dev/null +++ b/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", +] + diff --git a/simpleautodiff/simpleautodiff.py b/simpleautodiff/simpleautodiff.py index eba2668..3ba034d 100644 --- a/simpleautodiff/simpleautodiff.py +++ b/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))) + ) +