def trace_grad(fn, args):
  """Trace a function, and return a VJP and the function's output."""
  from tensorflow.python.eager.backprop import make_vjp
  result, vjp = make_vjp(fn)(*args)
  return result, vjp
Trace a function, and return a VJP and the function's output.