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