reverse_mode.py

import inspect
from functools import wraps
import ast
import textwrap
import math

# Example Node class for variable nodes
BinOp_LIBRARY = {
    ast.Add: ((lambda x,y: x+y), (lambda x,y, z_prime: (z_prime, z_prime))),
    ast.Sub: ((lambda x,y: x-y), (lambda x,y, z_prime: (z_prime, -z_prime))),
    ast.Mult: ((lambda x,y: x*y), (lambda x,y, z_prime: (z_prime*y, z_prime*x))),
    ast.Div: ((lambda x,y: x/y), (lambda x,y, z_prime: (z_prime/y, -z_prime*x/y**2))),
    ast.Pow: ((lambda x,y: x**y), (lambda x,y, z_prime: (
        z_prime*y*x**(y-1), x**y*math.log(x) if x>0 else 0
    ))),
    "exp": ((lambda a,x: n**x), (lambda a,x, z_prime: z_prime*z**x*math.log(a)))
}

TRIGON_LIBRARY = {
    "math.sin": ((lambda x: math.sin(x)), (lambda x, z_prime: (z_prime*math.cos(x),))),
    "math.cos": ((lambda x: math.cos(x)), (lambda x, z_prime: (-z_prime*math.sin(x),))),
    # "los"
}
# Transformer for modifying constants and function calls

def get_func_name(fs):
    try:
        f_index = list(TRIGON_LIBRARY.values()).index(fs)
        return list(TRIGON_LIBRARY.keys())[f_index]
    except ValueError as e:
        f_index = list(BinOp_LIBRARY.values()).index(fs)
        return list(BinOp_LIBRARY.keys())[f_index]

class GraphGenerator(ast.NodeTransformer):
    def __init__(self, args):
        self.__input_count = len(args)
        self.graph = [0]*len(args)
        self.input_args = args
        self.redundent_node_trace = {}
        self.name_nodes = {}
        self.__outputs = [0]*len(args)

    
    def prepare(self):
        # self.graph = tuple(self.graph)
        # self.reverse_graph = tuple(reversed(self.graph))


        self.__enum_graph = tuple(enumerate(self.graph))
        self.__enum_reverse_graph = tuple(enumerate(tuple(reversed(self.graph))))
        self.__graph_length = len(self.graph)
        self.__primes = [0]*self.__graph_length

    def visit_Constant(self, node):
        if(nd:=self.redundent_node_trace.get(node.value)):
            return nd

        node.trace_id = len(self.graph)
        self.graph.append(1)
        self.redundent_node_trace[node.value] = node
        self.__outputs.append(node.value)
        print(f"cnst: {node.value}, {len(self.graph)-1}, {id(node)}")
        print()

        return node

    def visit_name_manual(self, node):
        if (nd:=self.name_nodes.get(node.id)):
            return nd
        else:
            self.name_nodes[node.id] = node
            node.trace_id = self.input_args.index(node.id)


        print(f"var: {node.id}, {node.trace_id} , {id(node)}")
        print()
        return node

    def visit_BinOp(self, node):

        if isinstance(node.left, ast.Name):
            node.left = self.visit_name_manual(node.left)

        if isinstance(node.right, ast.Name):
            node.right = self.visit_name_manual(node.right)


        self.generic_visit(node)
        if (nd:=self.redundent_node_trace.get((type(node.op), node.left.trace_id, node.right.trace_id ))):
            return nd
        if (
            isinstance(node.op, (ast.Add, ast.Mult)) 
            and (nd:=self.redundent_node_trace.get((type(node.op), node.right.trace_id, node.left.trace_id )))
             
        ):
            return nd
        push_forward, pullback = BinOp_LIBRARY[type(node.op)]

        node.trace_id = len(self.graph)
        graph_node = (push_forward, pullback, (node.left.trace_id, node.right.trace_id))
        self.graph.append(graph_node)
        self.redundent_node_trace[(type(node.op), node.left.trace_id, node.right.trace_id)] = node
        self.__outputs.append(0)

        print(f"{node.left.trace_id}, {id(node.left)}")
        print(f"{node.right.trace_id}, {id(node.right)}")
        print(f"{type(node.op)}, {len(self.graph)-1}, {id(node)}")
        print()
        return node


    def visit_Call(self, node):
        self.generic_visit(node)


        call_string = node.func.value.id+'.'+node.func.attr \
            if isinstance(node.func, ast.Attribute) else node.func.id
        
        # CHECKS IF THE FUNCTION CALL IS CORRECT
        if not (isinstance(node.func, ast.Attribute) and node.func.value.id == 'math'
                        and call_string in TRIGON_LIBRARY.keys()):
            raise AttributeError(
                f"any function call must be from {TRIGON_LIBRARY.keys()}. not '{call_string}'"
            )

        # ONE argument per function call
        if len(node.args) != 1:
            raise ValueError(f"'{call_string}' takes 1 argument given {len(node.args)}")
        # single_arg = node.args[0]

        if isinstance(node.args[0], ast.Name):
            node.args[0] = self.visit_name_manual(node.args[0])
        
        self.generic_visit(node)
        if (nd:=self.redundent_node_trace.get((call_string, node.args[0].trace_id ))):
            return nd
        
        push_forward, pullback = TRIGON_LIBRARY[call_string]

        node.trace_id = len(self.graph)
        graph_node = (push_forward, pullback, (node.args[0].trace_id,))
        self.graph.append(graph_node)
        self.redundent_node_trace[(call_string, node.args[0].trace_id)] = node
        self.__outputs.append(0)

        print(f"{node.args[0].trace_id}, {id(node.args[0])}")
        print(f"{call_string}, {len(self.graph)-1}, --> {id(node)}")
        print()
        return node


    def gradient(self, *args):
        enum_graph = self.__enum_graph
        enum_reverse_graph = self.__enum_reverse_graph
        graph_length = self.__graph_length
        outputs = self.__outputs
        primes = self.__primes


        for i in range(self.__input_count):
            outputs[i] = args[i]

        for i, graph_node in enum_graph:
            if graph_node == 0:
                outputs[i] = args[i]
                continue
            if graph_node == 1:
                continue
            push_forward, _, indeces = graph_node
            inputs = [outputs[index] for index in indeces]
            outputs[i] = push_forward(*inputs)

        # backward pass
        primes = [0]*graph_length
        primes[-1] = 1
        gradients = [0]*self.__input_count
        for i, graph_node in enum_reverse_graph:
            idx = graph_length-i-1
            if graph_node == 0:
                gradients[idx]= primes[idx]
                # primes[idx]=0.0
                continue
            if graph_node == 1:
                # primes[idx] = 0.0
                continue

            _, pullback, indeces = graph_node
            inputs = [outputs[index] for index in indeces]
            
            try:
                local_primes = pullback(*inputs, primes[graph_length-i-1])
            except ValueError as e:
                print("debug error")
                print(*inputs, math.log(inputs[0]), graph_length-i-1)
                exit()
            # primes[graph_length-i-1] = 0.0
            for index, prime in zip(indeces, local_primes):
                primes[index] += prime

        return gradients

    def __call__(self, *args):
        try:
            return self.gradient(*args)
        except AttributeError as e:
            raise(e)



class GeneratorHelper(ast.NodeVisitor):
    def __init__(self):
        self.arg_nodes = []
        self.constants = []
    def visit_Return(self, node):
        print(self.arg_nodes)
        print(self.constants)
        print("#########")
        generator = GraphGenerator(self.arg_nodes)
        generator.visit(node)
        self.generator = generator
    def visit_arg(self, node):
        self.arg_nodes.append(node.arg)
    

from pprint import pprint
# Decorator to modify function behavior
def create_graph(func):

    source_code = textwrap.dedent(inspect.getsource(func))
    tree = ast.parse(source_code)
    graph = []
    # print(tree.body)
    generator_helper = GeneratorHelper()
    generator_helper.visit(tree)

    generator_helper.generator.prepare()
    return generator_helper.generator
    @wraps(func)
    def wrapper(*args, **kwargs):
        return modified_func(*args, **kwargs)

    return wrapper