[functorch] fix functionalize(): properly propagate input mutations (pytorch/functorch#654)

Author: bdhirsh

functorch/functorch/_src/eager_transforms.py

@@ -1224,10 +1224,10 @@ def wrapped(*args, **kwargs):
             func_args = _wrap_all_tensors_to_functional(args, func_level)
             func_kwargs = _wrap_all_tensors_to_functional(kwargs, func_level)
 
-            flattened_unwrapped_args = tree_flatten(args)
-            flattened_wrapped_args = tree_flatten(func_args)
-            flattened_unwrapped_kwargs = tree_flatten(kwargs)
-            flattened_wrapped_kwargs = tree_flatten(func_kwargs)
+            flattened_unwrapped_args, _ = tree_flatten(args)
+            flattened_wrapped_args, _ = tree_flatten(func_args)
+            flattened_unwrapped_kwargs, _ = tree_flatten(kwargs)
+            flattened_wrapped_kwargs, _ = tree_flatten(func_kwargs)
 
             func_outputs = func(*func_args, **func_kwargs)
             outputs = _unwrap_all_tensors_from_functional(func_outputs)

functorch/test/test_eager_transforms.py

@@ -2776,6 +2776,9 @@ def _check_functionalize_correctness(self, f, inpt):
         # Check that outputs are the same
         self.assertEqual(actual_outputs, expected_outputs)
 
+        # Inputs might have been mutated by f: check that they were mutated properly
+        self.assertEqual(inpt1, inpt2)
+
     def test_simple_view(self, device):
         def f(x: torch.Tensor) -> torch.Tensor:
             tmp = torch.ones(2, device=device)