def test_resolve():
def g(x):
return 2 * x
def f(x):
return g(x)
node = annotate.resolve_calls(f)
assert anno.getanno(node.body[0].body[0].value, 'func') == g
def f(x):
return h(x)
node = quoting.parse_function(f)
with pytest.raises(AttributeError):
annotate.resolve_calls(f)
def tangent(f):
"""A decorator which removes the `with grad_of` statement.
This allows the function to be called as usual.
Args:
f: A function
Returns:
A function with any `with grad_of` context managers removed.
"""
node = annotate.resolve_calls(f)
RemoveWith().visit(node)
wrapped = functools.wraps(f)(compile_.compile_function(node))
wrapped.tangent = f
return wrapped
def autodiff_ast(func, wrt, motion, mode, preserve_result, check_dims,
verbose):
"""Perform AD on a single function and return the AST.
Args:
See `grad`.
Returns:
node: The AST of a module containing the adjoint and primal function
definitions.
required: A list of non-built in functions that this function called, and
of which the primals and adjoints need to be made available in order
for the returned function to run.
"""
node = annotate.resolve_calls(func)
node = desugar.explicit_loop_indexes(node)
fence.validate(node, inspect.getsource(func))
node = anf_.anf(node)
if verbose >= 2:
print('ANF')
print(quoting.to_source(node))
if mode == 'reverse':
node, required, stack = reverse_ad.reverse_ad(node.body[0], wrt,
preserve_result,
check_dims)
if verbose >= 2:
print('RAW')
print(quoting.to_source(node))
if motion == 'split':
node = reverse_ad.split(node, stack)
else:
node = reverse_ad.joint(node)
if verbose >= 2:
print('MOTION')
print(quoting.to_source(node))
elif mode == 'forward':
node, required = forward_ad.forward_ad(node.body[0], wrt,
preserve_result, check_dims)
return node, required
