vmap Basics

Jump to bottom

Manuel edited this page Jul 3, 2024 · 3 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!

pFad - Phonifier reborn

Pfad - The Proxy pFad of © 2024 Garber Painting. All rights reserved.

Note: This service is not intended for secure transactions such as banking, social media, email, or purchasing. Use at your own risk. We assume no liability whatsoever for broken pages.

Alternative Proxies:

Alternative Proxy