def tangent(f):
  """A decorator which removes the `with insert_grad_of` statement.

  This allows the function to be called as usual.

  Args:
    f: A function

  Returns:
    A function with any `with insert_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