Integrate `functorch` or wait for `torch.vmap` in `PyTorch`

[Zymrael]

Several methods and classes in torchdyn would benefit from vjp/jvp: LNN, MSForward and others. We should decide on a way to integrate the functional API of PyTorch and then merge into torchdyn the prototypes using the internal vmap API.

[zou3519]

Hi, I'm one of the authors of functorch. We now have binary builds of functorch (e.g. one can do pip install functorch) so including functorch as a dependency is much easier now :)

[Zymrael]

Great news @zou3519, we will begin integrating functorch in several torchdyn classes.

[zou3519]

Please let us know if you run into any problems (feel free to open issues over at https://github.com/pytorch/functorch), functorch is actively being developed and we'd be interested in helping support your use case.