


x = torch.randn([5])
def f(x, scale=4.):
  return x * scale
(_, vjpfunc) = torch.func.vjp(f, x)
vjps = vjpfunc(torch.ones_like(x))
assert torch.allclose(vjps[0], torch.full(x.shape, 4.))



Note:
 Using PyTorch "torch.no_grad" together with "vjp". Case 1: Using
 "torch.no_grad" inside a function:

 >>> def f(x):
 >>>     with torch.no_grad():
 >>>         c = x ** 2
 >>>     return x - c

 In this case, "vjp(f)(x)" will respect the inner
 "torch.no_grad".Case 2: Using "vjp" inside "torch.no_grad"
 context manager:

 >>> with torch.no_grad():
 >>>     vjp(f)(x)

 In this case, "vjp" will respect the inner "torch.no_grad", but
 not the outer one. This is because "vjp" is a "function
 transform": its result should not depend on the result of a
 context manager outside of "f".
