pytorch · bohnstingl · Feb 1, 2025 · Feb 3, 2025 · Feb 6, 2025 · Feb 7, 2025
diff --git a/test/functorch/test_control_flow.py b/test/functorch/test_control_flow.py
@@ -3486,18 +3486,20 @@ def f(fct, init, xs):
             """\
 def forward(self, fct_1, init_1, xs_1):
     permute = torch.ops.aten.permute.default(xs_1, [0, 1, 2])
-    select_copy = torch.ops.aten.select_copy.int(permute, 0, 0)
+    flip = torch.ops.aten.flip.default(permute, [0]);  permute = None
+    select_copy = torch.ops.aten.select_copy.int(flip, 0, 0)
     add = torch.ops.aten.add.Tensor(init_1, select_copy);  add = None
     add_1 = torch.ops.aten.add.Tensor(init_1, select_copy);  select_copy = add_1 = None
     sym_size_int_1 = torch.ops.aten.sym_size.int(init_1, 1)
     sym_size_int_2 = torch.ops.aten.sym_size.int(init_1, 2)
     sym_size_int_3 = torch.ops.aten.sym_size.int(xs_1, 1)
     sym_size_int_4 = torch.ops.aten.sym_size.int(xs_1, 2);  xs_1 = None
     scan_combine_graph_0 = self.scan_combine_graph_0
-    scan = torch.ops.higher_order.scan(scan_combine_graph_0, [init_1], [permute], True, [sym_size_int_1, sym_size_int_2, sym_size_int_3, sym_size_int_4]);  scan_combine_graph_0 = init_1 = permute = sym_size_int_1 = sym_size_int_2 = sym_size_int_3 = sym_size_int_4 = None
+    scan = torch.ops.higher_order.scan(scan_combine_graph_0, [init_1], [flip], [sym_size_int_1, sym_size_int_2, sym_size_int_3, sym_size_int_4]);  scan_combine_graph_0 = init_1 = flip = sym_size_int_1 = sym_size_int_2 = sym_size_int_3 = sym_size_int_4 = None
     getitem = scan[0]
     getitem_1 = scan[1];  scan = None
-    return (getitem, getitem_1)""",  # noqa: B950
+    flip_1 = torch.ops.aten.flip.default(getitem_1, [0]);  getitem_1 = None
+    return (getitem, flip_1)""",  # noqa: B950
         )
 
         # Check graph
@@ -3516,10 +3518,11 @@ def forward(self, L_init_ : torch.Tensor, L_xs_ : torch.Tensor):
     v = l_init_ + select_copy;  v = None
     x = l_init_ + select_copy;  select_copy = x = None
     scan_combine_fn_0 = self.scan_combine_fn_0
-    scan = torch.ops.higher_order.scan(scan_combine_fn_0, [l_init_], [elem], True, []);  scan_combine_fn_0 = l_init_ = elem = None
+    scan = torch.ops.higher_order.scan(scan_combine_fn_0, [l_init_], [elem_1], []);  scan_combine_fn_0 = l_init_ = elem_1 = None
     getitem = scan[0]
-    getitem_1 = scan[1];  scan = None
-    return (getitem, getitem_1)""",  # noqa: B950
+    elem_2 = scan[1];  scan = None
+    flip_1 = torch.flip(elem_2, [0]);  elem_2 = None
+    return (getitem, flip_1)""",  # noqa: B950
         )
 
     @skipIfNoDynamoSupport
@@ -7484,7 +7487,7 @@ def forward(self, L_init_ : torch.Tensor, L_xs_ : torch.Tensor, L_add_closure_0_
     r_4 = r_3.add(l_add_closure_0_cell_contents_1_0_);  r_3 = None
     r_5 = r_4.sum();  r_4 = r_5 = None
     scan_combine_fn_0 = self.scan_combine_fn_0
-    scan = torch.ops.higher_order.scan(scan_combine_fn_0, [l_init_], [r], False, [l_add_closure_0_cell_contents_0_param_, l_add_closure_0_cell_contents_1_0_]);  scan_combine_fn_0 = l_init_ = r = l_add_closure_0_cell_contents_0_param_ = l_add_closure_0_cell_contents_1_0_ = None
+    scan = torch.ops.higher_order.scan(scan_combine_fn_0, [l_init_], [r], [l_add_closure_0_cell_contents_0_param_, l_add_closure_0_cell_contents_1_0_]);  scan_combine_fn_0 = l_init_ = r = l_add_closure_0_cell_contents_0_param_ = l_add_closure_0_cell_contents_1_0_ = None
     getitem = scan[0]
     getitem_1 = scan[1];  scan = None
    return (getitem, getitem_1)""",  # noqa: B950
@@ -7505,7 +7508,7 @@ def forward(self, L_init_ : torch.Tensor, L_xs_ : torch.Tensor, L_add_closure_0_
     ret = matmul_1 + l_add_closure_0_cell_contents_1_0_;  matmul_1 = None
     x = ret.sum();  ret = x = None
     scan_combine_fn_0 = self.scan_combine_fn_0
-    scan = torch.ops.higher_order.scan(scan_combine_fn_0, [l_init_], [elem], False, [l_add_closure_0_cell_contents_0_param_, l_add_closure_0_cell_contents_1_0_]);  scan_combine_fn_0 = l_init_ = elem = l_add_closure_0_cell_contents_0_param_ = l_add_closure_0_cell_contents_1_0_ = None
+    scan = torch.ops.higher_order.scan(scan_combine_fn_0, [l_init_], [elem], [l_add_closure_0_cell_contents_0_param_, l_add_closure_0_cell_contents_1_0_]);  scan_combine_fn_0 = l_init_ = elem = l_add_closure_0_cell_contents_0_param_ = l_add_closure_0_cell_contents_1_0_ = None
     getitem = scan[0]
     getitem_1 = scan[1];  scan = None
     return (getitem, getitem_1)""",  # noqa: B950

diff --git a/torch/_dynamo/variables/higher_order_ops.py b/torch/_dynamo/variables/higher_order_ops.py
@@ -1429,12 +1429,10 @@ def call_function(
 
         args, kwargs = LazyVariableTracker.realize_all((args, kwargs))
 
-        def arg_extractor(combine_fn, init, xs, reverse, additional_inputs):
-            return combine_fn, init, xs, reverse, additional_inputs
+        def arg_extractor(combine_fn, init, xs, additional_inputs):
+            return combine_fn, init, xs, additional_inputs
 
-        combine_fn, init, xs, reverse, additional_inputs = arg_extractor(
-            *args, **kwargs
-        )
+        combine_fn, init, xs, additional_inputs = arg_extractor(*args, **k
8000
wargs)
         assert isinstance(additional_inputs, variables.BaseListVariable)
 
         if xs.python_type() != list:
@@ -1542,7 +1540,6 @@ def _check_phs_position_match(
             init_proxy,
             xs_proxy,
             # dim.as_proxy(),
-            reverse.as_proxy(),
             additional_inputs_proxy,
         )
 

@@ -325,6 +325,9 @@ def add(x: torch.Tensor, y: torch.Tensor):
     for elem in leaves_xs_orig:
         leaves_xs.append(torch.movedim(elem, dim, 0))
 
+    if reverse:
+        leaves_xs = [torch.flip(elem, [0]) for elem in leaves_xs]
+
     out = combine_fn(
         pytree.tree_unflatten(leaves_init, spec_init),
         pytree.tree_unflatten([first_slice_copy(elem) for elem in leaves_xs], spec_xs),
@@ -393,10 +396,8 @@ def _check_new_carry_match_init(leaves_init, leaves_carry):
         num_inp_leaves=len(leaves_xs),
     )
 
-    def run_flattened_scan(combine_fn, leaves_init, leaves_xs, reverse):
-        return scan_op(
-            combine_fn, leaves_init, leaves_xs, reverse, additional_inputs=[]
-        )
+    def run_flattened_scan(combine_fn, leaves_init, leaves_xs):
+        return scan_op(combine_fn, leaves_init, leaves_xs, additional_inputs=[])
 
     if not torch._dynamo.is_compiling():
         from torch._dynamo.backends.debugging import (
@@ -415,16 +416,18 @@ def run_flattened_scan(combine_fn, leaves_init, leaves_xs, reverse):
                     combine_fn,
                     leaves_init,
                     leaves_xs,
-                    reverse=reverse,
                 )
     else:
-        result = run_flattened_scan(combine_fn, leaves_init, leaves_xs, reverse)
+        result = run_flattened_scan(combine_fn, leaves_init, leaves_xs)
 
     result_carry, result_flat = _extract_carry_and_out(
         result,
         len(leaves_init),
     )
 
+    if reverse:
+        result_flat = [torch.flip(elem, [0]) for elem in result_flat]
+
     return pytree.tree_unflatten(result_carry, spec_init), pytree.tree_unflatten(
         result_flat, tree_out
     )
@@ -434,16 +437,16 @@ class ScanOp(HigherOrderOperator):
     def __init__(self):
         super().__init__("scan")
 
-    def __call__(self, combine_fn, init, xs, reverse, additional_inputs):
+    def __call__(self, combine_fn, init, xs, additional_inputs):
         assert isinstance(additional_inputs, list), "additional_inputs must be a list."
         validate_subgraph_args_types(additional_inputs)
-        return super().__call__(combine_fn, init, xs, reverse, additional_inputs)
+        return super().__call__(combine_fn, init, xs, additional_inputs)
 
 
 scan_op = ScanOp()
 
 
-def generic_scan(operator, init, xs, dim=0, reverse=False, additional_inputs=None):
+def generic_scan(operator, init, xs, dim=0, additional_inputs=None):
     additional_inputs = additional_inputs if additional_inputs is not None else []
 
     def _scan(init, xs):
@@ -453,10 +456,7 @@ def _scan(init, xs):
             return carry, []
 
         num_elems = xs[0].shape[dim]
-        if reverse:
-            ind = num_elems - 1
-        else:
-            ind = 0
+        ind = 0
 
         # Compute dummy shapes for the pre-allocation
         num_init_leaves = len(init)
@@ -497,7 +497,7 @@ def store_out_in_outs(out, ind):
                 o.scatter_(0, ind * idx, x.unsqueeze(0))
 
         for i in range(num_elems):
-            ind = i if not reverse else num_elems - i - 1
+            ind = i
             carry, out = _extract_carry_and_out(
                 operator(
                     *carry,
@@ -537,7 +537,6 @@ def trace_scan(
     combine_fn: Callable,
     init: list[torch.Tensor],
     xs: list[torch.Tensor],
-    reverse: bool,
     additional_inputs: list[torch.Tensor],
 ):
     from torch._dynamo.utils import clone_input
@@ -582,7 +581,7 @@ def trace_scan(
 
     proxy_mode.tracer.root.register_module(combine_graph_name, combine_graph)
 
-    args = (combine_graph, init, xs, reverse, additional_inputs)
+    args = (combine_graph, init, xs, additional_inputs)
     proxy_args = pytree.tree_map(proxy_mode.tracer.unwrap_proxy, args)
     out_proxy = proxy_mode.tracer.create_proxy(
         "call_function", func_overload, proxy_args, {}, name="scan"
@@ -602,12 +601,10 @@ def trace_scan(
 
 
 @scan_op.py_impl(DispatchKey.CompositeExplicitAutograd)
-def scan_op_dense(combine_fn, init, xs, reverse, additional_inputs):
+def scan_op_dense(combine_fn, init, xs, additional_inputs):
     mode = _get_current_dispatch_mode()
     assert mode is None, "Mode should never be enabled for CPU/CUDA key"
-    return generic_scan(
-        combine_fn, init, xs, reverse=reverse, additional_inputs=additional_inputs
-    )
+    return generic_scan(combine_fn, init, xs, additional_inputs=additional_inputs)
 
 
 class ScanAutogradOp(torch.autograd.Function):
@@ -896,12 +893,12 @@ def scan_autograd(combine_fn, init, xs, reverse, additional_inputs):
 
 
 @scan_op.py_impl(ProxyTorchDispatchMode)
-def scan_proxy_mode(mode, combine_fn, init, xs, reverse, additional_inputs):
-    return trace_scan(mode, scan_op, combine_fn, init, xs, reverse, additional_inputs)
+def scan_proxy_mode(mode, combine_fn, init, xs, additional_inputs):
+    return trace_scan(mode, scan_op, combine_fn, init, xs, additional_inputs)
 
 
 @scan_op.py_impl(FakeTensorMode)
-def scan_fake_tensor_mode(mode, combine_fn, init, xs, reverse, additional_inputs):
+def scan_fake_tensor_mode(mode, combine_fn, init, xs, additional_inputs):
     with mode:
         scan_length = xs[0].shape[0]
         carry, outputs = _extract_carry_and_out(
@@ -920,7 +917,7 @@ def scan_fake_tensor_mode(mode, combine_fn, init, xs, reverse, additional_inputs
 
 
 @scan_op.py_functionalize_impl
-def scan_functionalize(ctx, combine_fn, init, xs, reverse, additional_inputs):
+def scan_functionalize(ctx, combine_fn, init, xs, additional_inputs):
     unwrapped_xs = ctx.unwrap_tensors(xs)
     unwrapped_init = ctx.unwrap_tensors(init)
     unwrapped_additional_inputs = ctx.unwrap_tensors(additional_inputs)
@@ -951,7 +948,6 @@ def scan_functionalize(ctx, combine_fn, init, xs, reverse, additional_inputs):
             functional_combine_fn,
             unwrapped_init,
             unwrapped_xs,
-            reverse,
             unwrapped_additional_inputs,
         )
     return ctx.wrap_tensors(ret)