Repository: pytorch
Issue: Dynamo capture for HigherOrderOperators, followups
Labels: triaged,tracker,oncall: pt2,module: dynamo,module: higher order operators

General higher order operators tracking issue. @voznesenskym, @Chillee 

##  General UX

- [x] calling nn.Modules or using nn.Module parameters should not graph break
- [x] HigherOrderOperator body functions should not all be named "cond_body"
- [x] body functions with non-single Tensor output
- [x] https://github.com/pytorch/pytorch/pull/103870
- [x] HigherOrderOperator fallback should fallback on the entire HigherOrderOperator instead of letting Dynamo introspect `HigherOrderOperator.__call__`. Likely needs some sort of bytecode rewriting.
- [x] HigherOrderOperator side effect handling (e.g. https://github.com/pytorch/pytorch/issues/103613)
- [x] https://github.com/pytorch/pytorch/issues/104298
- [ ] body functions with non-Tensor inputs
- [ ] HigherOrderOperators should accept pytree input/output
- [ ] Constants should probably not be inputs to the body function
- [ ] Unification: torch.utils.checkpoint and torch.func.grad use a slightly different mechanism for lifting the APIs into a HigherOrderOperator. Ideally we would have a single mechanism.


## autograd.Function
- Nothing here

## checkpoint
- [ ] Revisit why we need to add guards for constants - see ctx here - https://github.com/pytorch/pytorch/pull/102448

## Control flow operators
- [ ] Get cond, map to use the new mechanism, delete the old mechanism
- [ ] Decide on the stable {cond, map, scan, while_loop, for_i_loop} APIs. Fix existing APIs (cond, map) to match what we want.
- [ ] Turn Dynamo on by default for control flow operators in eager-mode PyTorch.
- [ ] differentiable cond
- [ ] capturable while-loop
- [ ] Integrate Brian's AOTExport API to stop doing recomputation in the backwards pass.

## functorch transforms
### grad
- [ ] func which takes kwargs is not supported (currently graph-break) : Fix (speculate_subgraph should support kwargs)
- [ ] Implement support for `torch.func.grad_and_value` (we would get `grad` for free that)
- [ ] speculate_subgraph should ideally support pytrees
- [ ] functools.wraps nested function use leads to graph-break (https://github.com/pytorch/pytorch/issues/103913)
- [ ] side-effect handling in dynamo is incomplete/buggy (mainly related to how it handles closure), fixing this would should help us avoid modelling grad mode for `torch.func.grad` in `speculate_subgraph` (see also https://github.com/pytorch/pytorch/pull/102264#discussion_r1231112556)
- [ ] figure out why we need to disable(func) to actually get this to work. disable(func) has the side effect of disabling func when it is used outside of grad, which is undesirable
- [ ] Some mutations are not treated as side-effects leading to incorrect graph (see `test_escaped_wrappers_are_marked_as_dead`). Also, ref https://github.com/pytorch/pytorch/issues/103613
- [ ] Lifted tensor are treated as side-effect (see `test_grad_freevar_tensor`)

## # vmap
- [ ] API to allow tracing through `torch` functions (which don't deal with tensors) https://github.com/pytorch/pytorch/pull/101707#discussion_r1283429266

cc @ezyang @bdhirsh @anijain2305 @chauhang @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @chenyang78 @kadeng @ydwu4 @msaroufim @aakhundov @soumith @wconstab @ngimel @Xia-Weiwen @desertfire