def run_and_trace(fun, x, *args, **kwargs):
start_node = VJPNode.new_root()
start_box = new_box(x, 0, start_node)
out = fun(start_box, *args, **kwargs)
return start_box, out
def calc_jacobian_elem(start, end):
with warnings.catch_warnings():
warnings.simplefilter("ignore")
b = anp.ones(start.shape)
n = new_box(b, 0, VJPNode.new_root())
jac = backward_pass(n, end._node)
return jac._value
def _autograd_is_indep_analytic(func, *args, **kwargs):
"""Test analytically whether a function is independent of its arguments
using Autograd.
Args:
func (callable): Function to test for independence
args (tuple): Arguments for the function with respect to which
to test for independence
kwargs (dict): Keyword arguments for the function at which
(but not with respect to which) to test for independence
Returns:
bool: Whether the function seems to not depend on it ``args``
analytically. That is, an output of ``True`` means that the
``args`` do *not* feed into the output.
In Autograd, we test this by sending a ``Box`` through the function and
testing whether the output is again a ``Box`` and on the same trace as
the input ``Box``. This means that we can trace actual *independence*
of the output from the input, not only whether the passed function is
constant.
The code is adapted from
`autograd.tracer.py::trace
<https://github.com/HIPS/autograd/blob/master/autograd/tracer.py#L7>`__.
"""
# pylint: disable=protected-access
node = VJPNode.new_root()
with trace_stack.new_trace() as t:
start_box = new_box(args, t, node)
end_box = func(*start_box, **kwargs)
if type(end_box) in [tuple, list]:
if any(
isbox(_end) and _end._trace == start_box._trace
for _end in end_box):
return False
elif isinstance(end_box, np.ndarray):
if end_box.ndim == 0:
end_box = [end_box.item()]
if any(
isbox(_end) and _end._trace == start_box._trace
for _end in end_box):
return False
else:
if isbox(end_box) and end_box._trace == start_box._trace:
return False
return True
def calc_jacobian(start, end):
# if the end_box is not a box - autograd can not track back
if not isbox(end):
return vspace(start.shape).zeros()
# the final jacobian matrices
jac = []
# the backward pass is done for each objective function once
for j in range(end.shape[1]):
b = anp.zeros(end.shape)
b[:, j] = 1
n = new_box(b, 0, VJPNode.new_root())
_jac = backward_pass(n, end._node)
jac.append(_jac)
jac = anp.stack(jac, axis=1)
return jac
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 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_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 autograd_box(x):
"""Box a tensor in AutoGrad."""
t = trace_stack.new_trace().__enter__()
n = VJPNode.new_root()
return new_box(x, t, n)