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

DiffEqML / torchdyn

A PyTorch library entirely dedicated to neural differential equations, implicit models and related numerical methods

https://torchdyn.org

Apache License 2.0

1.36k stars 125 forks source link

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

Open Zymrael opened 3 years ago

Zymrael commented 3 years ago

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 commented 2 years ago

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 commented 2 years ago

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

zou3519 commented 2 years ago

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.