def create_forward(out_node):
  """Create a user-friendly forward function.

  Ensures that a single value instead of a tuple is returned if the user asked
  for the gradient with respect to only one input.

  Args:
    out_node: The function definition AST.

  Returns:
    The function definition with potentially changed return statement.
  """
  retval = out_node.body[0].body[-1]
  if len(retval.value.elts) == 1:
    retval.value = retval.value.elts[0]
  return out_node