import torch


def jvp(y, x, v, create_graph=False):
    """Jacobian-vector product dy/dx @ v.
    Pass `create_graph=True` to allow backprop on output.
    """
    return torch.autograd.grad(y.flatten(), x, v.flatten(),
        retain_graph=True,
        create_graph=create_graph,
        only_inputs =True)[0]