code: def _create_forward(out_node):
    
    retval = out_node.body[0].body[-1]
    if len(retval.value.elts) == 1:
      retval.value = retval.value.elts[0]
    return out_node
docs: 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.