vmap Basics

Jump to bottom

Richard Zou edited this page Jul 25, 2025 · 4 revisions

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

I would love to contribute to PyTorch!