


expectedX = torch.diag(torch.ones_like(x))
expectedY = torch.diag(2 * y)
assert torch.allclose(jacobian[0], expectedX)
assert torch.allclose(jacobian[1], expectedY)



Note:
 Using PyTorch "torch.no_grad" together with "jacrev". 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, "jacrev(f)(x)" will respect the inner
 "torch.no_grad".Case 2: Using "jacrev" inside "torch.no_grad"
 context manager:

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

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