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 trace(fun, start_nodes, args):
with tracer.trace_stack.new_trace() as t:
start_boxes = [
tracer.new_box(x, t, n) for x, n in zip(args, start_nodes)
]
end_box = fun(*start_boxes)
if tracer.isbox(end_box) and end_box._trace == t:
return end_box._value, end_box._node
else:
warnings.warn("Output seems independent of input.")
return end_box, None
def test_value_and_grad():
fun = lambda x: np.sum(np.sin(x)**2)
dfun = grad(fun)
dfun_both = value_and_grad(fun)
x = npr.randn(5)
assert not isbox(dfun_both(x)[0])
check_equivalent(fun(x), dfun_both(x)[0])
check_equivalent(dfun(x), dfun_both(x)[1])
def fun2(x): return dfun_both(x)[0]
check_grads(fun2)(x)
def test_value_and_grad():
fun = lambda x: np.sum(np.sin(x)**2)
dfun = grad(fun)
dfun_both = value_and_grad(fun)
x = npr.randn(5)
assert not isbox(dfun_both(x)[0])
check_equivalent(fun(x), dfun_both(x)[0])
check_equivalent(dfun(x), dfun_both(x)[1])
def fun2(x):
return dfun_both(x)[0]
check_grads(fun2)(x)
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 is_constant(x):
return not ag_tracer.isbox(x)
def backtrack(output_boxes, input_node_map):
"""Trace the computation graph from output boxes to ancestors in the input.
Given the output of a function called on boxed inputs, backtrack finds
which, if any, of the nodes in the input_node_map each output element
depends on.
Parameters
----------
output_boxes: list
List of outputs from a traced function.
input_node_map: dict
Dictionary of TracerNode -> box elements representing the potential root
nodes of the computation graph.
Returns
-------
output_dependencies: dict
Dictionary mapping output_idx to a (deduplicated) list of nodes
in the input that output idx depends on. For instance, if we have
def f(x, y, z):
return x + y, y + z
Then, we get output_dependencies {0: [Node_x, Node_y], 1: [Node_y, Node_z]}.
"""
if not isinstance(output_boxes, Iterable):
output_boxes = [output_boxes]
# For each output, figure out which (if any) of the inputs it depends on
# by solving a graph search problem on the computation graph.
output_dependencies = dict(
(output_idx, set()) for output_idx in range(len(output_boxes)))
for idx, output_box in enumerate(output_boxes):
# If the output isn't a box, then it's independent of all the inputs.
if isbox(output_box):
# Use breadth-first search to find parents
# pylint: disable-msg=protected-access
queue = [output_box._node]
while queue:
node = queue.pop(0)
# Check if ancestor is an input node
if node in input_node_map:
input_box = input_node_map[node]
if output_box._trace == input_box._trace:
output_dependencies[idx].add(node)
# If the node corresponds to an irrelevant dependency,
# e.g. 0 * constant, we skip it.
if is_dead_node(node):
continue
# Add the parents to the queue
for parent in node.parents:
# This is a hack to handle a common pattern where the
# user returns z = np.array([x1, x2]) from a function and
# then we inspect one item, e.g. z[0]. By default, both x1 and x2
# are parents of z, but in this case, we can detect that
# z[0] only depends on x1. In general, however, if after
# constructing z, we apply another function to the array,
# e.g. f(z) we cannot detect the subset of x that the result depends on.
# This is one of the limitations of using boxes based on autograd.
# Note this limitation means that we might add extra
# edges to the causal graph. It will never mean we miss an
# edge, so the returned graphs are still valid.
if is_getitem_node(node) and is_array_node(parent):
index = node.args[1]
queue.append(parent.parents[index])
else:
queue.append(parent)
return dict((k, list(v)) for k, v in output_dependencies.items())
