rzou 99257002fa Extend auto_functionalized to support ops that return Tensors (#115135) ... We can auto-functionalize operators that mutate their inputs as long as the outputs of the operator do not alias their inputs. The user needs to provide an abstract impl for the operator if it has non-trivial returns. - We update can_auto_functionalize(op) to include ops that return (but do not alias) Tensors - We update auto_functionalized(op, mutated_args_names, kwargs) to return (out, mutated_args), where `out = op(**kwargs)` and `mutated_args` are the new values of the inputs that would have been mutated. Test Plan: - new test Pull Request resolved: https://github.com/pytorch/pytorch/pull/115135 Approved by: https://github.com/bdhirsh ghstack dependencies: #114955, #114956, #115134		2023-12-05 22:43:06 +00:00
..
__init__.py	Add Fake Cross Ref Mode, migrate sparse to it (#85382)	2022-09-21 17:15:47 +00:00
fake_tensor.py	Refactor can_auto_functionalize (#115134)	2023-12-05 22:43:06 +00:00
fake_utils.py	Expose Flash attn to autograd (#114378)	2023-12-01 23:42:06 +00:00
functional_tensor.py	Extend auto_functionalized to support ops that return Tensors (#115135)	2023-12-05 22:43:06 +00:00
meta_utils.py	Expand dynamic dims support for traceable subclasses (#114311)	2023-12-05 21:09:25 +00:00
schema_check_mode.py	Use pytree.tree_leaves everywhere (#112324)	2023-10-30 03:39:04 +00:00