Things that would be nice to have with the functorch release

[zou3519]

The functorch release will be approximately 3/1. Below is a list of issues that would be good to resolve, ordered by priority.

P0

Theme: important bugs and testing

• Calling VMap on Grad - Issues with Indexing #381
• vmapvjp failure on __getitem__ #393
• Printing tensor broken in vmap(grad( case #382
• Optimization: convolution_backward doesn't always need to call .contiguous on certain inputs pytorch#70184
• @zou3519 Incorrect results when involving .detach() with second-order derivatives #353 (warning: might be tricky)
• @samdow loud error message for autograd.function and functorch in error case (autograd.Function is silently incorrect on functorch transforms #207)
• @zou3519 Port NTK example to functorch
• @zou3519 Internal assert in NTK computation #417
• @samdow Make sure the examples still work
• @zou3519 Some more docs and tutorials #360
• Port https://github.com/srush/annotated-s4 to functorch
• Evaluation on torchbenchmark models. E.g. (1) grad transform works on those, (2) compute per-sample-grads using vmap+grad, ...
• make_functional doesn't work with parameter sharing #446
• index_put doesn't have a batch rule #124
• functorch build+test time went from 50m to 20m in the last two months #454

P0.5

Theme: Other fixes and testing

• Optimize convolution batching rule performance  #365
• Confirm that we do have tests for no-batch-dim support and the batching rules actually support them #443
• svd-related op regression in functorch #428
• max_pool1d regression in functorch #424
• has_aux flag for jacfwd (and maybe hessian) #392
• Harden {jvp, jacfwd, hessian} #342 (Just the performance testing parts of these, especially with jvp and jacfwd)
• functorch.jvp doesn't work with AOTAutograd #437
• Audit docs website
• @samdow Make batch norm error message more informative + write up how to fix models that use batch norm and run under vmap

P1

It's OK if these don't make it into the release (we can do a follow-up later). But they are very nice to have.

• @samdow Initial vmap randomness design
• fix batch_norm in omniglot example #384
• Ban calling specific autograd APIs from inside a function transform #376
• Switch make_functional over to nn.utils._stateless and/or clarify semantics (Update make_functional after https://github.com/pytorch/pytorch/pull/71542 #415)
• Version selector for the docs #451