vmap Basics - pytorch/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

Next: Try out the vmap Onboarding Lab

Unit 8: function transforms/Training Loops (Optional) - vmap-onboarding-lab