def forward_ad(node, wrt, preserve_result=False, check_dims=True):
  """Perform forward-mode AD on an AST.

  This function analyses the AST to determine which variables are active and
  proceeds by taking the naive derivative. Before returning the primal and
  adjoint it annotates push and pop statements as such.

  Args:
    node: A `FunctionDef` AST node.
    wrt: A tuple of argument indices with respect to which we take the
        derivative.
    preserve_result: A boolean indicating whether the original
        non-differentiated function value should be returned
    check_dims: A boolean indicating whether the provided derivatives should
        have the same shape as their corresponding arguments.

  Returns:
    mod: A `Module` node containing the naive primal and adjoint of the
        function which can be fed to the `split` and `joint` functions.
    required: A list of tuples of functions and argument indices. These
        functions were called by the function but did not have an adjoint.
  """
  if not isinstance(node, gast.FunctionDef):
    raise TypeError

  # Activity analysis
  cfg_obj = cfg.CFG.build_cfg(node)
  cfg.Active(range(len(node.args.args))).visit(cfg_obj.entry)

  # Build forward mode function
  fad = ForwardAD(wrt, preserve_result, check_dims)
  node = fad.visit(node)

  # Annotate stacks
  node = annotate.find_stacks(node)

  # Clean up naive forward-mode fcode
  node = gast.Module([node])
  anno.clearanno(node)

  return node, fad.required
Perform forward-mode AD on an AST.

  This function analyses the AST to determine which variables are active and
  proceeds by taking the naive derivative. Before returning the primal and
  adjoint it annotates push and pop statements as such.

  Args:
    node: A `FunctionDef` AST node.
    wrt: A tuple of argument indices with respect to which we take the
        derivative.
    preserve_result: A boolean indicating whether the original
        non-differentiated function value should be returned
    check_dims: A boolean indicating whether the provided derivatives should
        have the same shape as their corresponding arguments.

  Returns:
    mod: A `Module` node containing the naive primal and adjoint of the
        function which can be fed to the `split` and `joint` functions.
    required: A list of tuples of functions and argument indices. These
        functions were called by the function but did not have an adjoint.