Functorch

Jax-like function transformations for PyTorch neural networks!

The main idea is that you can separate object-based PyTorch modules into function transformations consisting of a params sequence and a apply function.

Can also vmap.

Resonance

  • https://wandb.ai/functorch-examples/functorch-examples/reports/Working-with-FuncTorch-An-Introduction–VmlldzoxNzMxNDI1
  • https://github.com/soumik12345/functorch-examples/blob/main/01_classfier_training.ipynb