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
def _assert_tangent_parse_error(func, fragment):
try:
fence.validate(quoting.parse_function(func), inspect.getsource(func))
assert False
except fence.TangentParseError as expected:
assert fragment in str(expected)