Update on "auto functionalize base_hop"

cc ezyang SherlockNoMad EikanWang jgong5 wenzhe-nrv voznesenskym penguinwu Guobing-Chen XiaobingSuper zhuhaozhe blzheng jiayisunx ipiszy chenyang78 kadeng muchulee8 amjames chauhang aakhundov [ghstack-poisoned]
pytorch · ydwu4 · Apr 10, 2025 · Apr 11, 2025 · Apr 11, 2025 · Apr 12, 2025
commit 2c27c6e2a68003e3f4c66e1c15f165ef5ea83621
diff --git a/torch/_higher_order_ops/utils.py b/torch/_higher_order_ops/utils.py
@@ -908,7 +908,7 @@ def flat_fn(*args_and_grad_outs):
         tangents = args_and_grad_outs[n_primals:]
         grad_args = bw_fn(primals, tangents)[1]
         assert len(args) == len(grad_args)
-        out_tuple = tuple(
+        out = tuple(
             (
                 torch.zeros_like(arg)
                 if isinstance(arg, torch.Tensor) and grad is None
@@ -919,9 +919,6 @@ def flat_fn(*args_and_grad_outs):
         _maybe_clone = clone_outputs_aliasing_inputs(args_and_grad_outs)
         return pytree.tree_map(_maybe_clone, out)
 
-        maybe_clone = clone_outputs_aliasing_inputs(args_and_grad_outs)
-        return pytree.tree_map(maybe_clone, out_tuple)
-
     return flat_fn