Now add the pseudo hack in the emitter

Author: JacobSzwejbka

exir/emit/_emitter.py

@@ -1189,6 +1189,19 @@ def _emit_control_flow(
         """
         specs = self.node.meta["spec"]
 
+        # For scan, set the shape_dynamism for the stacked outputs (y_outputs) to DYNAMIC_BOUND
+        # BEFORE emitting the specs. This is because et_copy_index has cat shape semantics but
+        # stack memory behavior, so we need to be able to update the shape +1 for each iteration
+        # which we can't do for tensors marked static.
+        if target is torch.ops.higher_order.scan:
+            combine_fn, init, xs, additional_inputs = args
+            num_carry = len(init)
+            if isinstance(specs, (list, tuple)):
+                y_specs = specs[num_carry:]
+                for y_spec in y_specs:
+                    if isinstance(y_spec, TensorSpec):
+                        y_spec.shape_dynamism = TensorShapeDynamism.DYNAMIC_BOUND
+
         subemitter_binding_output_values = pytree.tree_map(
             lambda spec: self._emit_spec(spec),
             specs,

exir/pass_base.py

@@ -572,7 +572,7 @@ def call_scan(
             xs_element_data.append(ph.meta["val"])
 
         combine_fn_result = self.call_submodule(
-            combine_fn, (*init , *xs_element_data , *additional_inputs)
+            combine_fn, (*init, *xs_element_data, *additional_inputs)
         )
         assert combine_fn_result is not None
 

exir/passes/spec_prop_pass.py

@@ -59,11 +59,13 @@ def __call__(self, graph_module: torch.fx.GraphModule) -> PassResult:
         res = ExportPass()(graph_module)
         assert res is not None
         gm = res.graph_module
+
         def get_spec(x):
             if hasattr(x, "meta"):
                 return x.meta.get("spec", None)
             else:
                 return None
+
         for module in gm.modules():
             if isinstance(module, torch.fx.GraphModule):
                 for node in module.graph.nodes: