vmap basics - subramen/pytorch GitHub Wiki

Page Maintainers: @zou3519

Scope

  • Understand what vmap is
  • Understand how vmap is implemented with Tensor Subclassing
  • Understand (as a PyTorch developer) how to add missing support for an operation with vmap

Introduction to vmap

Read through the functorch whirlwind tour.

How to add vmap support for a PyTorch operator

Read through https://github.com/pytorch/pytorch/blob/master/functorch/writing_batching_rules.md

Try out the vmap Onboarding Lab

vmap-onboarding-lab

PyTorch presented to you with love by the PyTorch Team of contributors

⚠️ **GitHub.com Fallback** ⚠️