(v)
| 97 | # Using `vmap`, we can vectorize the whole computation, computing the |
| 98 | # Jacobian in a single call to `autograd.grad`. |
| 99 | def get_vjp(v): |
| 100 | return torch.autograd.grad(y, x, v)[0] |
| 101 | |
| 102 | jacobian_vmap = vmap(get_vjp)(basis_vectors) |
| 103 | assert torch.allclose(jacobian_vmap, jacobian) |
nothing calls this directly
no outgoing calls
no test coverage detected