Future Feature Request: Vectorization

[Parskatt]

Pytorch is (slowly) introducing vmap: https://pytorch.org/docs/master/generated/torch.vmap.html?highlight=vmap#torch.vmap
When this feature becomes stable it seems like a great addition for jacobian calculation, probably giving additional performance.

[sgould]

This is cool. Thanks for recommending. Once vmap makes it into an official release we'll look to profile and use it.