I need to compute both the vector-Jacobian product and the Jacobian-vector product at the same time, and then to backprop through both. I have the following code that I have tested and I believe it works correctly:
def vjp(f, x, v, create_graph=True):
    x = x.detach().requires_grad_()
    y = f(x)
    y.backward(v, create_graph=create_graph)
    return x.grad

def jvp(f, x, v, create_graph=True):
    g = lambda v: vjp(f, x, v, create_graph=True)
    return vjp(g, v, v, create_graph=create_graph)

def get_loss(f, x, v):
    vjp_val = vjp(f, x, v)
    jvp_val = jvp(f, x, v)

    return (vjp_val - jvp_val).norm(1)

It is however inefficient, as it effectively computes f(x).backward(v) twice. Hence I would like to rewrite it in such way that it only does so once. Here is my attempt:
def get_loss_fast(f, x, v):
    x = x.detach().requires_grad_()
    y = f(x)
    y.backward(v, create_graph=True)
    vjp_val = x.grad

    vjp_val.backward(v, create_graph=True)
    jvp_val = x.grad

    return (vjp_val - jvp_val).norm(1)

This code always returns zero. In fact, inside the get_loss_fast function, vjp_val is jvp_val is True, which means that the second backward() does not overwrite the output of the first one.
How can I compute this loss efficiently and correctly?