# Path: autograd/tracer.py
# def trace(start_node, fun, x):
#     def __init__(self, value, fun, args, kwargs, parent_argnums, parents):
#     def initialize_root(self, *args, **kwargs):
#     def new_root(cls, *args, **kwargs):
# def primitive(f_raw):
#     def f_wrapped(*args, **kwargs):
# def register_notrace(trace_type, primitive_fun):
# def notrace_primitive(f_raw):
#     def f_wrapped(*args, **kwargs):
# def find_top_boxed_args(args):
#     def __init__(self):
#     def new_trace(self):
#     def __init__(self, value, trace, node):
#     def __bool__(self):
#     def __str__(self):
#     def register(cls, value_type):
# def new_box(value, trace, node):
# class Node(object):
# class TraceStack(object):
# class Box(object):
# 
# Path: autograd/util.py
# def func(f): return f
# 
# def subval(x, i, v):
#     x_ = list(x)
#     x_[i] = v
#     return tuple(x_)


from itertools import count
from functools import reduce
from .tracer import trace, primitive, toposort, Node, Box, isbox, getval
from .util import func, subval
import warnings

# -------------------- reverse mode --------------------

def make_vjp(fun, x):
    start_node = VJPNode.new_root()
    end_value, end_node =  trace(start_node, fun, x)
    if end_node is None:
        def vjp(g): return vspace(x).zeros()
    else:
        def vjp(g): return backward_pass(g, end_node)
    return vjp, end_value

def backward_pass(g, end_node):
    outgrads = {end_node : (g, False)}
    for node in toposort(end_node):
        outgrad = outgrads.pop(node)
        ingrads = node.vjp(outgrad[0])
        for parent, ingrad in zip(node.parents, ingrads):
            outgrads[parent] = add_outgrads(outgrads.get(parent), ingrad)
    return outgrad[0]


# Path: autograd/tracer.py
# def trace(start_node, fun, x):
#     def __init__(self, value, fun, args, kwargs, parent_argnums, parents):
#     def initialize_root(self, *args, **kwargs):
#     def new_root(cls, *args, **kwargs):
# def primitive(f_raw):
#     def f_wrapped(*args, **kwargs):
# def register_notrace(trace_type, primitive_fun):
# def notrace_primitive(f_raw):
#     def f_wrapped(*args, **kwargs):
# def find_top_boxed_args(args):
#     def __init__(self):
#     def new_trace(self):
#     def __init__(self, value, trace, node):
#     def __bool__(self):
#     def __str__(self):
#     def register(cls, value_type):
# def new_box(value, trace, node):
# class Node(object):
# class TraceStack(object):
# class Box(object):
# 
# Path: autograd/util.py
# def func(f): return f
# 
# def subval(x, i, v):
#     x_ = list(x)
#     x_[i] = v
#     return tuple(x_)


# Path: autograd/core.py
from itertools import count
from functools import reduce
from .tracer import trace, primitive, toposort, Node, Box, isbox, getval
from .util import func, subval
import warnings


# -------------------- reverse mode --------------------

def make_vjp(fun, x):
    start_node = VJPNode.new_root()
    end_value, end_node =  trace(start_node, fun, x)
    if end_node is None:
        def vjp(g): return vspace(x).zeros()
    else:
        def vjp(g): return backward_pass(g, end_node)
    return vjp, end_value

def backward_pass(g, end_node):
    outgrads = {end_node : (g, False)}
    for node in toposort(end_node):
        outgrad = outgrads.pop(node)
        ingrads = node.vjp(outgrad[0])
        for parent, ingrad in zip(node.parents, ingrads):
            outgrads[parent] = add_outgrads(outgrads.get(parent), ingrad)
    return outgrad[0]


class VJPNode(Node):