


assert torch.allclose(vjps[0], x.cos() + -x.sin())



"vjp()" can even support outputs being Python structs



x = torch.randn([5])
f = lambda x: {'first': x.sin(), 'second': x.cos()}
(_, vjpfunc) = torch.func.vjp(f, x)
cotangents = {'first': torch.ones([5]), 'second': torch.ones([5])}
vjps = vjpfunc(cotangents)
assert torch.allclose(vjps[0], x.cos() + -x.sin())



The function returned by "vjp()" will compute the partials with
   respect to each of the "primals"



x, y = torch.randn([5, 4]), torch.randn([4, 5])
(_, vjpfunc) = torch.func.vjp(torch.matmul, x, y)
cotangents = torch.randn([5, 5])
vjps = vjpfunc(cotangents)
assert len(vjps) == 2
assert torch.allclose(vjps[0], torch.matmul(cotangents, y.transpose(0, 1)))
assert torch.allclose(vjps[1], torch.matmul(x.transpose(0, 1), cotangents))



"primals" are the positional arguments for "f". All kwargs use
   their default value



x = torch.randn([5])


