def jacfun(*args, **kwargs):
start_node, end_nodes, tape = forward_pass(list_fun, args, kwargs,
argnum)
grads = list(
map(partial(backward_pass, start_node, tape=tape), end_nodes))
shape = dummy.outshape + getshape(args[argnum])
return np.reshape(concatenate(grads), shape) if shape else grads[0]
def jacfun(*args, **kwargs):
start_node, end_nodes, tape = forward_pass(list_fun, args, kwargs,
argnum)
run = partial(backward_pass, start_node, tape=tape, preserve_tape=True)
grads = [run(end_node) for end_node in end_nodes]
del tape[:]
shape = dummy.outshape + getshape(args[argnum])
return np.reshape(concatenate(grads), shape) if shape else grads[0]
def time_fan_out_fan_in_forward_pass():
if MASTER_BRANCH:
forward_pass(fan_out_fan_in, (2.,), {})
else:
start_node = VJPNode.new_root()
trace(start_node, fan_out_fan_in, x)
def time_long_forward_pass():
if MASTER_BRANCH:
forward_pass(f_long, (2.,), {})
else:
start_node = VJPNode.new_root()
trace(start_node, f_long, x)
def time_exp_call():
onp.exp(2.)
def time_exp_primitive_call_unboxed():
np.exp(2.)
def time_exp_primitive_call_boxed():
if MASTER_BRANCH:
np.exp(progenitor)
else:
np.exp(start_box)
def time_no_autograd_control():
# Test whether the benchmarking machine is running slowly independent of autograd
A = np.random.randn(200, 200)
np.dot(A, A)
if MASTER_BRANCH:
short_start_node, short_end_node = forward_pass(f_short, (2.,), {})
long_start_node, long_end_node = forward_pass(f_long, (2.,), {})
fan_start_node, fan_end_node = forward_pass(fan_out_fan_in, (2.,), {})
progenitor = new_progenitor(2.)
else:
x = 2.
start_node = VJPNode.new_root()
start_box = new_box(x, 0, start_node)
_, short_end_node = trace(VJPNode.new_root(), f_short, x)
_, long_end_node = trace(VJPNode.new_root(), f_long, x)
_, fan_end_node = trace(VJPNode.new_root(), fan_out_fan_in, x)
def jacfun(*args, **kwargs):
start_node, end_nodes, tape = forward_pass(list_fun, args, kwargs, argnum)
grads = list(map(partial(backward_pass, start_node, tape=tape), end_nodes))
shape = dummy.outshape + getshape(args[argnum])
return np.reshape(concatenate(grads), shape) if shape else grads[0]
def time_long_forward_pass():
core.forward_pass(f_long, (2., ), {})
def time_fan_out_fan_in_forward_pass():
core.forward_pass(fan_out_fan_in, (2., ), {})
def time_fan_out_fan_in_forward_pass():
if MASTER_BRANCH:
forward_pass(fan_out_fan_in, (2.,), {})
else:
start_node = VJPNode.new_root(x)
trace(start_node, fan_out_fan_in, x)
def time_short_forward_pass():
core.forward_pass(f_short, (2., ), {})
def time_long_forward_pass():
if MASTER_BRANCH:
forward_pass(f_long, (2.,), {})
else:
start_node = VJPNode.new_root(x)
trace(start_node, f_long, x)
def time_exp_call():
onp.exp(2.)
def time_exp_primitive_call_unboxed():
np.exp(2.)
def time_exp_primitive_call_boxed():
if MASTER_BRANCH:
np.exp(progenitor)
else:
np.exp(start_box)
def time_no_autograd_control():
# Test whether the benchmarking machine is running slowly independent of autograd
A = np.random.randn(200, 200)
np.dot(A, A)
if MASTER_BRANCH:
short_start_node, short_end_node = forward_pass(f_short, (2.,), {})
long_start_node, long_end_node = forward_pass(f_long, (2.,), {})
fan_start_node, fan_end_node = forward_pass(fan_out_fan_in, (2.,), {})
progenitor = new_progenitor(2.)
else:
x = 2.
start_node = VJPNode.new_root(x)
start_box = new_box(x, 0, start_node)
_, short_end_node = trace(VJPNode.new_root(x), f_short, x)
_, long_end_node = trace(VJPNode.new_root(x), f_long, x)
_, fan_end_node = trace(VJPNode.new_root(x), fan_out_fan_in, x)
def time_short_forward_pass():
core.forward_pass(f_short, (2.,), {})
def time_fan_out_fan_in_forward_pass():
core.forward_pass(fan_out_fan_in, (2.,), {})
def time_long_forward_pass():
core.forward_pass(f_long, (2.,), {})