File size: 3,951 Bytes
ab4488b |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 |
import warnings
from contextlib import contextmanager
from collections import defaultdict
from .util import subvals, toposort
from .wrap_util import wraps
def trace(start_node, fun, x):
with trace_stack.new_trace() as t:
start_box = new_box(x, t, start_node)
end_box = fun(start_box)
if isbox(end_box) and end_box._trace == start_box._trace:
return end_box._value, end_box._node
else:
warnings.warn("Output seems independent of input.")
return end_box, None
class Node(object):
__slots__ = []
def __init__(self, value, fun, args, kwargs, parent_argnums, parents):
assert False
def initialize_root(self, *args, **kwargs):
assert False
@classmethod
def new_root(cls, *args, **kwargs):
root = cls.__new__(cls)
root.initialize_root(*args, **kwargs)
return root
def primitive(f_raw):
"""
Wraps a function so that its gradient can be specified and its invocation
can be recorded. For examples, see the docs."""
@wraps(f_raw)
def f_wrapped(*args, **kwargs):
boxed_args, trace, node_constructor = find_top_boxed_args(args)
if boxed_args:
argvals = subvals(args, [(argnum, box._value) for argnum, box in boxed_args])
if f_wrapped in notrace_primitives[node_constructor]:
return f_wrapped(*argvals, **kwargs)
parents = tuple(box._node for _ , box in boxed_args)
argnums = tuple(argnum for argnum, _ in boxed_args)
ans = f_wrapped(*argvals, **kwargs)
node = node_constructor(ans, f_wrapped, argvals, kwargs, argnums, parents)
return new_box(ans, trace, node)
else:
return f_raw(*args, **kwargs)
f_wrapped.fun = f_raw
f_wrapped._is_autograd_primitive = True
return f_wrapped
notrace_primitives = defaultdict(set)
def register_notrace(trace_type, primitive_fun):
notrace_primitives[trace_type].add(primitive_fun)
def notrace_primitive(f_raw):
@wraps(f_raw)
def f_wrapped(*args, **kwargs):
argvals = map(getval, args)
return f_raw(*argvals, **kwargs)
f_wrapped._is_primitive = True
return f_wrapped
def find_top_boxed_args(args):
top_trace = -1
top_boxes = []
top_node_type = None
for argnum, arg in enumerate(args):
if isbox(arg):
trace = arg._trace
if trace > top_trace:
top_boxes = [(argnum, arg)]
top_trace = trace
top_node_type = type(arg._node)
elif trace == top_trace:
top_boxes.append((argnum, arg))
return top_boxes, top_trace, top_node_type
class TraceStack(object):
def __init__(self):
self.top = -1
@contextmanager
def new_trace(self):
self.top += 1
yield self.top
self.top -= 1
trace_stack = TraceStack()
class Box(object):
type_mappings = {}
types = set()
__slots__ = ['_value', '_trace', '_node']
def __init__(self, value, trace, node):
self._value = value
self._node = node
self._trace = trace
def __bool__(self):
return bool(self._value)
__nonzero__ = __bool__
def __str__(self):
return "Autograd {0} with value {1}".format(
type(self).__name__, str(self._value))
@classmethod
def register(cls, value_type):
Box.types.add(cls)
Box.type_mappings[value_type] = cls
Box.type_mappings[cls] = cls
box_type_mappings = Box.type_mappings
def new_box(value, trace, node):
try:
return box_type_mappings[type(value)](value, trace, node)
except KeyError:
raise TypeError("Can't differentiate w.r.t. type {}".format(type(value)))
box_types = Box.types
isbox = lambda x: type(x) in box_types # almost 3X faster than isinstance(x, Box)
getval = lambda x: getval(x._value) if isbox(x) else x
|