diff --git a/test/functorch/test_control_flow.py b/test/functorch/test_control_flow.py
index 621a640f374ff..aa3ea646c1ff3 100644
--- a/test/functorch/test_control_flow.py
+++ b/test/functorch/test_control_flow.py
@@ -129,7 +129,7 @@ def compile_mode_helper(fct, compile_mode):
         return fct
 
 
-def get_scan_combine_fn(name, associative=True):
+def get_scan_combine_fn(name, expand_with_carry=False, parameters=None):
     def add(x: torch.Tensor, y: torch.Tensor):
         return x + y
 
@@ -163,6 +163,18 @@ def non_pointwise(x: torch.Tensor, y: torch.Tensor):
         W = torch.diag(torch.ones(2, device=x.device))
         return x @ W + y @ W
 
+    def RNN(x: torch.Tensor, y: torch.Tensor):
+        c_new = y @ parameters[0] + parameters[1]
+        h_new = torch.tanh(c_new + x @ parameters[2] + parameters[3])
+        return h_new, h_new
+
+    def fct_c1_no_grad(x: torch.Tensor, y: torch.Tensor):
+        h_new = torch.tanh(x[0] + x[1] + y)
+        c2 = x[1] + y
+        with torch.no_grad():
+            c1 = x[0] + y
+        return (c1, c2), h_new
+
     if name == "add":
         fct = add
     elif name == "adds":
@@ -179,10 +191,14 @@ def non_pointwise(x: torch.Tensor, y: torch.Tensor):
         fct = complex_pointwise
     elif name == "non_pointwise":
         fct = non_pointwise
+    elif name == "RNN":
+        fct = RNN
+    elif name == "fct_c1_no_grad":
+        fct = fct_c1_no_grad
     else:
         raise ValueError("Combine_fn name unknown!")
 
-    if not associative:
+    if expand_with_carry:
         return lambda x, y: (fct(x, y), fct(x, y))
     else:
         return fct
@@ -1316,8 +1332,8 @@ def test_associative_scan_compile(
         scan_fct = compile_mode_helper(associative_scan, compile_mode)
 
         for op, op_pt in [
-            (get_scan_combine_fn("add", True), torch.cumsum),
-            (get_scan_combine_fn("mul", True), torch.cumprod),
+            (get_scan_combine_fn("add", False), torch.cumsum),
+            (get_scan_combine_fn("mul", False), torch.cumprod),
         ]:
             result = scan_fct(op, x, 0, reverse=reverse, combine_mode=combine_mode)
             result_exp = _fake_associative_scan(op, xs=x, dim=0, reverse=reverse)
@@ -1329,14 +1345,14 @@ def test_associative_scan_compile(
         # Jax Examples
         x = torch.arange(0, 4, device=device)
         cumsum1 = scan_fct(
-            get_scan_combine_fn("add", True),
+            get_scan_combine_fn("add", False),
             x,
             0,
             reverse=reverse,
             combine_mode=combine_mode,
         )
         cumsum_exp = _fake_associative_scan(
-            get_scan_combine_fn("add", True), x, 0, reverse=reverse
+            get_scan_combine_fn("add", False), x, 0, reverse=reverse
         )
         if not reverse:
             self.assertEqual(
@@ -1375,12 +1391,12 @@ def add2(x: torch.Tensor, y: torch.Tensor):
 
         for op, op_pt, init in [
             (
-                get_scan_combine_fn("add", False),
+                get_scan_combine_fn("add", True),
                 torch.cumsum,
                 torch.zeros(2, 2, device=device, requires_grad=autograd),
             ),
             (
-                get_scan_combine_fn("mul", False),
+                get_scan_combine_fn("mul", True),
                 torch.cumprod,
                 torch.ones(2, 2, device=device, requires_grad=autograd),
             ),
@@ -1399,14 +1415,14 @@ def add2(x: torch.Tensor, y: torch.Tensor):
         x = torch.arange(0, 4, device=device, dtype=torch.int64)
         init = torch.zeros(1, device=device, dtype=torch.int64)
         cumsum1 = scan_fct(
-            get_scan_combine_fn("add", False),
+            get_scan_combine_fn("add", True),
             init,
             x,
             dim=0,
             reverse=reverse,
         )
         cumsum_exp = _fake_scan(
-            get_scan_combine_fn("add", False),
+            get_scan_combine_fn("add", True),
             init=init,
             xs=x,
             dim=0,
@@ -1451,14 +1467,14 @@ def add2(x: torch.Tensor, y: torch.Tensor):
         )
         init = torch.ones(1, device=device, dtype=torch.float32, requires_grad=autograd)
         result = scan_fct(
-            get_scan_combine_fn("div", False),
+            get_scan_combine_fn("div", True),
             init,
             x,
             dim=0,
             reverse=reverse,
         )
         result_exp = _fake_scan(
-            get_scan_combine_fn("div", False),
+            get_scan_combine_fn("div", True),
             init=init,
             xs=x,
             dim=0,
@@ -1502,7 +1518,7 @@ def test_scan_dtype(self, reverse, compile_mode, device, dtype):
         # Check all outputs and carries on the correct device and with torch.float32
         x = torch.randn(3, 10, 2, device=device).to(dtype=dtype)
         op, init = (
-            get_scan_combine_fn("adds"),
+            get_scan_combine_fn("adds", False),
             torch.zeros(10, 2, device=device, dtype=dtype),
         )
         result = scan_fct(op, init, x, dim=0, reverse=reverse)
@@ -1521,7 +1537,7 @@ def test_scan_dtype(self, reverse, compile_mode, device, dtype):
         # carry.dtype torch.float32 and output.dtype torch.float16
         x = torch.randn(3, 10, 2, device=device).to(dtype=dtype)
         op, init = (
-            get_scan_combine_fn("adds"),
+            get_scan_combine_fn("adds", False),
             torch.zeros(10, 2, device=device, dtype=torch.float32),
         )
         result = scan_fct(op, init, x, dim=0, reverse=reverse)
@@ -1539,7 +1555,7 @@ def test_scan_dtype(self, reverse, compile_mode, device, dtype):
         # carry.dtype torch.int64 and output.dtype torch.float32
         x = torch.randn(3, 10, 2, device=device)
         op, init = (
-            get_scan_combine_fn("adds"),
+            get_scan_combine_fn("adds", False),
             torch.zeros(10, 2, device=device, dtype=dtype),
         )
         result = scan_fct(op, init, x, dim=0, reverse=reverse)
@@ -1577,8 +1593,8 @@ def test_associative_scan_dim(self, combine_mode, reverse, device):
             x = torch.randn(*shapes, device=device)
 
             for op, op_pt in [
-                (get_scan_combine_fn("add", True), torch.cumsum),
-                (get_scan_combine_fn("mul", True), torch.cumprod),
+                (get_scan_combine_fn("add", False), torch.cumsum),
+                (get_scan_combine_fn("mul", False), torch.cumprod),
             ]:
                 result = associative_scan(
                     op, x, rnd_scan_dim, reverse=reverse, combine_mode=combine_mode
@@ -1612,12 +1628,12 @@ def test_scan_dim(self, reverse, compile_mode, device, autograd):
 
             for op, op_pt, init in [
                 (
-                    get_scan_combine_fn("add", False),
+                    get_scan_combine_fn("add", True),
                     torch.cumsum,
                     torch.zeros(*init_shapes, device=device, requires_grad=autograd),
                 ),
                 (
-                    get_scan_combine_fn("mul", False),
+                    get_scan_combine_fn("mul", True),
                     torch.cumprod,
                     torch.ones(*init_shapes, device=device, requires_grad=autograd),
                 ),
@@ -1663,14 +1679,14 @@ def test_associative_scan_binary_operator(self, combine_mode, reverse, device):
         elements = (A.repeat((timesteps, 1)), projected_inputs)
 
         result1 = associative_scan(
-            get_scan_combine_fn("s5_operator", True),
+            get_scan_combine_fn("s5_operator", False),
             elements,
             0,
             combine_mode=combine_mode,
             reverse=reverse,
         )
         expected_result = _fake_associative_scan(
-            get_scan_combine_fn("s5_operator", True), elements, 0, reverse=reverse
+            get_scan_combine_fn("s5_operator", False), elements, 0, reverse=reverse
         )
         self.assertEqual(
             result1,
@@ -1706,14 +1722,14 @@ def test_scan_binary_operator(self, reverse, device, autograd):
         )
 
         result = scan(
-            get_scan_combine_fn("s5_operator", False),
+            get_scan_combine_fn("s5_operator", True),
             init,
             elements,
             dim=0,
             reverse=reverse,
         )
         expected_result = _fake_scan(
-            get_scan_combine_fn("s5_operator", False),
+            get_scan_combine_fn("s5_operator", True),
             init=init,
             xs=elements,
             dim=0,
@@ -1757,14 +1773,14 @@ def test_associative_scan_tuple(self, combine_mode, reverse, device):
         inp = (x, y)
 
         result1 = associative_scan(
-            get_scan_combine_fn("tuple_fct", True),
+            get_scan_combine_fn("tuple_fct", False),
             inp,
             0,
             reverse=reverse,
             combine_mode=combine_mode,
         )
         expected_result = _fake_associative_scan(
-            get_scan_combine_fn("tuple_fct", True), inp, 0, reverse=reverse
+            get_scan_combine_fn("tuple_fct", False), inp, 0, reverse=reverse
         )
         self.assertEqual(result1, expected_result)
 
@@ -1779,14 +1795,14 @@ def test_scan_tuple(self, reverse, device):
         init = tuple(torch._ops.ops.aten.slice(e, 0, 0, 1, 1) for e in inp)
 
         result_same = scan(
-            get_scan_combine_fn("tuple_fct", False),
+            get_scan_combine_fn("tuple_fct", True),
             init,
             inp,
             dim=0,
             reverse=reverse,
         )
         expected_result = _fake_scan(
-            get_scan_combine_fn("tuple_fct", False),
+            get_scan_combine_fn("tuple_fct", True),
             init=init,
             xs=inp,
             dim=0,
@@ -1864,14 +1880,14 @@ def fct_pointwise(x, y):
         inp = {"i": x, "j": ([y], [{"o": z}])}
 
         result = associative_scan(
-            get_scan_combine_fn("complex_pointwise", True),
+            get_scan_combine_fn("complex_pointwise", False),
             inp,
             0,
             combine_mode=combine_mode,
             reverse=reverse,
         )
         expected_result = _fake_associative_scan(
-            get_scan_combine_fn("complex_pointwise", True), inp, 0, reverse=reverse
+            get_scan_combine_fn("complex_pointwise", False), inp, 0, reverse=reverse
         )
         self.assertEqual(result, expected_result)
 
@@ -1925,14 +1941,14 @@ def test_scan_complex_pytree(self, reverse, device):
         init = pytree.tree_unflatten(init_flat, inp_spec)
 
         result = scan(
-            get_scan_combine_fn("complex_pointwise", False),
+            get_scan_combine_fn("complex_pointwise", True),
             init,
             inp,
             dim=0,
             reverse=reverse,
         )
         expected_result = _fake_scan(
-            get_scan_combine_fn("complex_pointwise", False),
+            get_scan_combine_fn("complex_pointwise", True),
             init=init,
             xs=inp,
             dim=0,
@@ -1962,7 +1978,7 @@ def test_associative_scan_downstream_scan_matmul(
         def chain_fct(inp):
             W = torch.ones(2, 5, device=device)
             o = associative_scan(
-                get_scan_combine_fn("add", True),
+                get_scan_combine_fn("add", False),
                 inp,
                 1,
                 reverse=reverse,
@@ -1974,7 +1990,7 @@ def chain_fct(inp):
 
         inp = torch.randn(3, 10, 2, device=device)
         expected_result = _fake_associative_scan(
-            get_scan_combine_fn("add", True), inp, 1, reverse=reverse
+            get_scan_combine_fn("add", False), inp, 1, reverse=reverse
         ) @ torch.ones(2, 5, device=device)
         result1 = fct_cmp(inp)
         self.assertEqual(result1, expected_result)
@@ -2000,14 +2016,14 @@ def test_associative_scan_downstream_scan_scan(
         # Chain with scan
         def chain_fct_same_dim(inp):
             o1 = associative_scan(
-                get_scan_combine_fn("add", True),
+                get_scan_combine_fn("add", False),
                 inp,
                 1,
                 combine_mode=combine_mode,
                 reverse=reverse,
             )
             o2 = associative_scan(
-                get_scan_combine_fn("add", True),
+                get_scan_combine_fn("add", False),
                 o1,
                 1,
                 combine_mode=combine_mode,
@@ -2020,9 +2036,9 @@ def chain_fct_same_dim(inp):
         inp = torch.randn(3, 10, 2, device=device)
 
         expected_result = _fake_associative_scan(
-            get_scan_combine_fn("add", True),
+            get_scan_combine_fn("add", False),
             _fake_associative_scan(
-                get_scan_combine_fn("add", True), inp, 1, reverse=reverse
+                get_scan_combine_fn("add", False), inp, 1, reverse=reverse
             ),
             1,
             reverse=reverse,
@@ -2051,14 +2067,14 @@ def test_associative_scan_downstream_scan_scan_different_dim(
         # Chain with scan on different dim
         def chain_fct_different_dim(inp):
             o1 = associative_scan(
-                get_scan_combine_fn("add", True),
+                get_scan_combine_fn("add", False),
                 inp,
                 1,
                 combine_mode=combine_mode,
                 reverse=reverse,
             )
             o2 = associative_scan(
-                get_scan_combine_fn("add", True),
+                get_scan_combine_fn("add", False),
                 o1,
                 0,
                 combine_mode=combine_mode,
@@ -2070,9 +2086,9 @@ def chain_fct_different_dim(inp):
 
         inp = torch.randn(3, 10, 2, device=device)
         expected_result = _fake_associative_scan(
-            get_scan_combine_fn("add", True),
+            get_scan_combine_fn("add", False),
             _fake_associative_scan(
-                get_scan_combine_fn("add", True), inp, 1, reverse=reverse
+                get_scan_combine_fn("add", False), inp, 1, reverse=reverse
             ),
             0,
             reverse=reverse,
@@ -2095,7 +2111,7 @@ def test_scan_downstream_scan_matmul(self, compile_mode, reverse, device, autogr
             def chain_fct(inp):
                 W = torch.ones(2, 5, device=device)
                 o = scan(
-                    get_scan_combine_fn("add", False),
+                    get_scan_combine_fn("add", True),
                     init,
                     inp,
                     dim=1,
@@ -2106,7 +2122,7 @@ def chain_fct(inp):
             fct_cmp = compile_mode_helper(chain_fct, compile_mode)
 
             expected_result = _fake_scan(
-                get_scan_combine_fn("add", False),
+                get_scan_combine_fn("add", True),
                 init=init,
                 xs=inp,
                 dim=1,
@@ -2135,7 +2151,7 @@ def test_scan_downstream_scan_scan_dim(
 
         def chain_fct_different_dim(inp):
             o1 = scan(
-                get_scan_combine_fn("add", False),
+                get_scan_combine_fn("add", True),
                 init,
                 inp,
                 dim=1,
@@ -2143,7 +2159,7 @@ def chain_fct_different_dim(inp):
             )
             o1 = pytree.tree_map(lambda t: shift_source_dim_to_target_dim(t, 0, 1), o1)
             o2 = scan(
-                get_scan_combine_fn("add", False),
+                get_scan_combine_fn("add", True),
                 init2,
                 o1[1],
                 dim=0,
@@ -2154,7 +2170,7 @@ def chain_fct_different_dim(inp):
         fct_cmp = compile_mode_helper(chain_fct_different_dim, compile_mode)
 
         xs = _fake_scan(
-            get_scan_combine_fn("add", False),
+            get_scan_combine_fn("add", True),
             init=init,
             xs=inp,
             dim=1,
@@ -2162,7 +2178,7 @@ def chain_fct_different_dim(inp):
         )[1]
         xs = pytree.tree_map(lambda t: shift_source_dim_to_target_dim(t, 0, 1), xs)
         expected_result = _fake_scan(
-            get_scan_combine_fn("add", False),
+            get_scan_combine_fn("add", True),
             init=init2,
             xs=xs,
             dim=0,
@@ -2192,7 +2208,7 @@ def test_associative_scan_non_pointwise(self, reverse, device):
             "For combine_mode='pointwise', the combine_fn needs to be pointwise",
         ):
             out = associative_scan(
-                get_scan_combine_fn("non_pointwise", True),
+                get_scan_combine_fn("non_pointwise", False),
                 x,
                 0,
                 reverse=reverse,
@@ -2212,10 +2228,10 @@ def test_associative_scan_non_pointwise(self, reverse, device):
     def test_associative_scan_non_pointwise_generic(self, reverse, device):
         x = torch.randn(3, 10, 2, device=device)
         result_expected = _fake_associative_scan(
-            get_scan_combine_fn("non_pointwise", True), x, 0, reverse=reverse
+            get_scan_combine_fn("non_pointwise", False), x, 0, reverse=reverse
         )
         result1 = associative_scan(
-            get_scan_combine_fn("non_pointwise", True),
+            get_scan_combine_fn("non_pointwise", False),
             x,
             0,
             reverse=reverse,
@@ -2230,7 +2246,7 @@ def test_scan_non_pointwise(self, reverse, device):
         x = torch.randn(3, 10, 2, device=device)
         init = torch.randn(10, 2, device=device)
         result_expected = _fake_scan(
-            get_scan_combine_fn("non_pointwise", False),
+            get_scan_combine_fn("non_pointwise", True),
             init=init,
             xs=x,
             dim=0,
@@ -2238,7 +2254,7 @@ def test_scan_non_pointwise(self, reverse, device):
         )
 
         out = scan(
-            get_scan_combine_fn("non_pointwise", False),
+            get_scan_combine_fn("non_pointwise", True),
             init,
             x,
             dim=0,
@@ -2261,7 +2277,7 @@ def test_scan_compile_cnt(self, reverse, device):
             init = torch.randn(3, 5, device=device)
             # First compilation step
             torch.compile(scan, backend=cnt)(
-                get_scan_combine_fn("add", False),
+                get_scan_combine_fn("add", True),
                 init,
                 x,
                 dim=dim,
@@ -2273,7 +2289,7 @@ def test_scan_compile_cnt(self, reverse, device):
             init = torch.randn(3, 5, device=device)
             # Recompilation due to first different size
             torch.compile(scan, backend=cnt)(
-                get_scan_combine_fn("add", False),
+                get_scan_combine_fn("add", True),
                 init,
                 x,
                 dim=dim,
@@ -2285,7 +2301,7 @@ def test_scan_compile_cnt(self, reverse, device):
             init = torch.randn(3, 5, device=device)
             # No recompilation, because of dynamic shape
             torch.compile(scan, backend=cnt)(
-                get_scan_combine_fn("add", False),
+                get_scan_combine_fn("add", True),
                 init,
                 x,
                 dim=dim,
@@ -2297,7 +2313,7 @@ def test_scan_compile_cnt(self, reverse, device):
             init = torch.randn(3, 40, device=device)
             # Recompilation because of dim change
             torch.compile(scan, backend=cnt)(
-                get_scan_combine_fn("add", False),
+                get_scan_combine_fn("add", True),
                 init,
                 x,
                 dim=2,
@@ -2309,7 +2325,7 @@ def test_scan_compile_cnt(self, reverse, device):
             init = torch.randn(3, 40, device=device)
             # Recompilation due to first different size on new dim
             torch.compile(scan, backend=cnt)(
-                get_scan_combine_fn("add", False),
+                get_scan_combine_fn("add", True),
                 init,
                 x,
                 dim=2,
@@ -2321,7 +2337,7 @@ def test_scan_compile_cnt(self, reverse, device):
             init = torch.randn(3, 40, device=device)
             # No recompilation, because of dynamic shape on new dim
             torch.compile(scan, backend=cnt)(
-                get_scan_combine_fn("add", False),
+                get_scan_combine_fn("add", True),
                 init,
                 x,
                 dim=2,
@@ -2333,7 +2349,7 @@ def test_scan_compile_cnt(self, reverse, device):
             init = torch.randn(3, 40, device=device)
             # Recompilation because of dim change
             torch.compile(scan, backend=cnt)(
-                get_scan_combine_fn("add", False),
+                get_scan_combine_fn("add", True),
                 init,
                 x,
                 dim=1,
@@ -2345,7 +2361,7 @@ def test_scan_compile_cnt(self, reverse, device):
             init = torch.randn(3, 40, device=device)
             # Recompilation because of reverse change
             torch.compile(scan, backend=cnt)(
-                get_scan_combine_fn("add", False),
+                get_scan_combine_fn("add", True),
                 init,
                 x,
                 dim=1,
@@ -2357,7 +2373,7 @@ def test_scan_compile_cnt(self, reverse, device):
             init = torch.randn(3, 40, device=device)
             # No recompilation, as nothing changed
             torch.compile(scan, backend=cnt)(
-                get_scan_combine_fn("add", False),
+                get_scan_combine_fn("add", True),
                 init,
                 x,
                 dim=1,
@@ -2369,7 +2385,7 @@ def test_scan_compile_cnt(self, reverse, device):
             init = torch.randn(3, 80, device=device)
             # No recompilation, final test
             torch.compile(scan, backend=cnt)(
-                get_scan_combine_fn("add", False),
+                get_scan_combine_fn("add", True),
                 init,
                 x,
                 dim=1,
@@ -2399,7 +2415,7 @@ def test_scan_init_scanned_0(self, reverse, compile_mode, device):
             "Observed exception.*",
         ):
             result_init = scan_fct(
-                get_scan_combine_fn("add", False),
+                get_scan_combine_fn("add", True),
                 init,
                 inp,
                 dim=dim,
@@ -2425,7 +2441,7 @@ def test_scan_init_non_tensor(self, reverse, compile_mode, device):
             "Observed exception.*",
         ):
             result_init = scan_fct(
-                get_scan_combine_fn("add", False), init, x, dim=dim, reverse=reverse
+                get_scan_combine_fn("add", True), init, x, dim=dim, reverse=reverse
             )
 
     @unittest.skipIf(not SM70OrLater, "triton")
@@ -2450,7 +2466,7 @@ def test_scan_init_wrong_shape(self, reverse, compile_mode, device):
             "Observed exception.*",
         ):
             result_init = scan_fct(
-                get_scan_combine_fn("add", False),
+                get_scan_combine_fn("add", True),
                 init,
                 inp,
                 dim=dim,
@@ -2499,7 +2515,7 @@ def test_scan_init(self, reverse, compile_mode, device, autograd):
         # Only init and no input
         x = torch.randn(3, 1, 2, device=device, requires_grad=autograd)
         dim = 1
-        op, op_pt = (get_scan_combine_fn("add", False), torch.cumsum)
+        op, op_pt = (get_scan_combine_fn("add", True), torch.cumsum)
 
         # Only init given
         init = torch._ops.ops.aten.slice(x, dim, 0, 1, 1)
@@ -2516,7 +2532,7 @@ def test_scan_init(self, reverse, compile_mode, device, autograd):
         x = torch.randn(3, 5, 2, device=device, requires_grad=autograd)
         dim = 0
 
-        op, op_pt = (get_scan_combine_fn("add", False), torch.cumsum)
+        op, op_pt = (get_scan_combine_fn("add", True), torch.cumsum)
         inp = torch._ops.ops.aten.slice(x, dim, 1, None, 1)
 
         # Init tensor scalar
@@ -2580,7 +2596,7 @@ def add_scalar_carry_sliced_out(x: torch.Tensor, y: torch.Tensor):
             self.check_autograd(result_init, result_exp, (init, inp))
 
         # Correct case
-        op, op_pt = (get_scan_combine_fn("add", False), torch.cumsum)
+        op, op_pt = (get_scan_combine_fn("add", True), torch.cumsum)
         x = torch.randn(3, 2, 2, device=device, requires_grad=autograd)
         init = torch.zeros(3, 2, device=device, requires_grad=autograd)
         dim = 2
@@ -2677,7 +2693,7 @@ def test_scan_init_wrong_pytree_complex(self, reverse, device):
             ".*",
         ):
             result = scan(
-                get_scan_combine_fn("complex_pointwise", False),
+                get_scan_combine_fn("complex_pointwise", True),
                 init,
                 inp,
                 dim=0,
@@ -2763,14 +2779,14 @@ def fct_pointwise_different_carry(x, y):
             ),
         }
         result = scan_fct(
-            get_scan_combine_fn("complex_pointwise", False),
+            get_scan_combine_fn("complex_pointwise", True),
             init,
             inp,
             dim=0,
             reverse=reverse,
         )
         expected_result = _fake_scan(
-            get_scan_combine_fn("complex_pointwise", False),
+            get_scan_combine_fn("complex_pointwise", True),
             init,
             inp,
             dim=0,
@@ -2847,15 +2863,19 @@ def test_scan_closure_RNN(self, autograd):
             W_hh = W_hh.detach()
             b_hh = b_hh.detach()
 
-        def RNN(x: torch.Tensor, y: torch.Tensor):
-            c_new = y @ W_ih + b_ih
-            h_new = torch.tanh(c_new + x @ W_hh + b_hh)
-            return h_new, h_new
-
         expected_result = rnn(x, torch.unsqueeze(h, 0))
         expected_result_out = expected_result[0]
         expected_result_state = expected_result[1][0, :]
-        result = scan(RNN, h, x, dim=dim, reverse=False)
+
+        result = scan(
+            get_scan_combine_fn(
+                "RNN", expand_with_carry=False, parameters=[W_ih, b_ih, W_hh, b_hh]
+            ),
+            h,
+            x,
+            dim=dim,
+            reverse=False,
+        )
         result_cmp = [result[0], shift_source_dim_to_target_dim(result[1], 0, dim)]
         self.assertEqual(result_cmp[0], expected_result_state)
         self.assertEqual(result_cmp[1], expected_result_out)
@@ -2894,50 +2914,145 @@ def RNN(x: torch.Tensor, y: torch.Tensor):
     @unittest.skipIf(not SM70OrLater, "triton")
     @requires_cuda
     @parametrize("reverse", [False, True])
-    @parametrize("compile_mode", ["none", "eager", "compile", "compile_dynamic_shape"])
     @parametrize("device", [torch.device("cpu"), torch.device("cuda")])
-    def test_scan_closure_RNN_partial_autograd(self, reverse, compile_mode, device):
-        import random
+    @parametrize("autograd", [False, True])
+    def test_scan_closure_RNN_parameters_as_inputs(self, reverse, device, autograd):
+        x = torch.randn(3, 5, 10, device=device, requires_grad=autograd)
+        h = torch.randn(3, 7, device=device, requires_grad=autograd)
+        W_ih = torch.randn(5, 7, device=device, requires_grad=autograd)
+        b_ih = torch.randn(7, device=device, requires_grad=autograd)
+        W_hh = torch.randn(7, 7, device=device, requires_grad=autograd)
+        b_hh = torch.randn(7, device=device, requires_grad=autograd)
+
+        with self.assertRaisesRegex(
+            # Should be: RuntimeError,
+            # "All xs leaves must have a scan dimension > 0",
+            torch._dynamo.exc.Unsupported,
+            "Observed exception.*",
+        ):
+            result = scan(
+                get_scan_combine_fn(
+                    "RNN", expand_with_carry=False, parameters=[W_ih, b_ih, W_hh, b_hh]
+                ),
+                h,
+                [x, W_ih, b_ih, W_hh, b_hh],
+                dim=2,
+                reverse=reverse,
+            )
 
+    @unittest.skipIf(not SM70OrLater, "triton")
+    @requires_cuda
+    @parametrize("reverse", [False, True])
+    @parametrize("compile_mode", ["none", "eager", "compile", "compile_dynamic_shape"])
+    @parametrize(
+        "partial_grad", ["xs", "init", "additional_inputs", "complex", "random"]
+    )
+    @parametrize("device", [torch.device("cpu"), torch.device("cuda")])
+    def test_scan_closure_RNN_partial_autograd(
+        self, reverse, compile_mode, partial_grad, device
+    ):
         dim = 1
         scan_fct = compile_mode_helper(scan, compile_mode)
+
+        # The first two booleans are the xs
+        # The second two are the inits
+        # The last four are the additional_inputs
         autograds = []
-        autograds.append([True, True, True, True, True, True])
-        autograds.append([False, False, True, True, True, True])
-        autograds.append([True, True, False, False, False, False])
-        autograds.append([True, False, False, False, False, False])
-        autograds.append([False, True, False, False, False, False])
-        for _ in range(5):
-            autograds.append([bool(random.randint(0, 1)) for _ in range(6)])
+
+        if partial_grad == "xs":
+            # xs tests
+            autograds.append([True, False, True, True, True, True, True, True])
+            autograds.append([False, False, True, True, True, True, True, True])
+        elif partial_grad == "init":
+            # init tests
+            autograds.append([True, True, False, True, True, True, True, True])
+            autograds.append([True, True, False, False, True, True, True, True])
+        elif partial_grad == "additional_inputs":
+            # additional input tests
+            autograds.append([True, True, True, True, False, True, False, True])
+            autograds.append([True, True, True, True, False, False, False, False])
+        elif partial_grad == "complex":
+            # complex cases
+            autograds.append([True, False, False, False, False, False, False, True])
+            autograds.append([False, False, True, True, False, False, False, True])
+        elif partial_grad == "random":
+            # random tests
+            import random
+
+            for _ in range(5):
+                autograds.append([bool(random.randint(0, 1)) for _ in range(8)])
 
         for autograd in autograds:
             x = torch.randn(3, 10, 5, device=device, requires_grad=autograd[0])
-            h = torch.randn(3, 7, device=device, requires_grad=autograd[1])
-            W_ih = torch.randn(5, 7, device=device, requires_grad=autograd[2])
-            b_ih = torch.randn(7, device=device, requires_grad=autograd[3])
-            W_hh = torch.randn(7, 7, device=device, requires_grad=autograd[4])
-            b_hh = torch.randn(7, device=device, requires_grad=autograd[5])
-
-            params = [p for p, a in zip([x, h, W_ih, b_ih, W_hh, b_hh], autograd) if a]
+            x1 = torch.randn(3, 10, 5, device=device, requires_grad=autograd[1])
+            h = torch.randn(3, 7, device=device, requires_grad=autograd[2])
+            h_1 = torch.randn(3, 7, device=device, requires_grad=autograd[3])
+            W_ih = torch.randn(5, 7, device=device, requires_grad=autograd[4])
+            b_ih = torch.randn(7, device=device, requires_grad=autograd[5])
+            W_hh = torch.randn(7, 7, device=device, requires_grad=autograd[6])
+            b_hh = torch.randn(7, device=device, requires_grad=autograd[7])
+
+            params = [
+                p
+                for p, a in zip([x, x1, h, h_1, W_ih, b_ih, W_hh, b_hh], autograd)
+                if a
+            ]
 
             def RNN(x: torch.Tensor, y: torch.Tensor):
-                c_new = y @ W_ih + b_ih
-                h_new = torch.tanh(c_new + x @ W_hh + b_hh)
-                return c_new, h_new
-
-            if False in autograd:
-                with self.assertRaisesRegex(
-                    RuntimeError,
-                    "scan currently only supports Autograd if all.*",
-                ):
-                    result = scan_fct(RNN, h, x, dim=dim, reverse=reverse)
-            else:
-                result = scan_fct(RNN, h, x, dim=dim, reverse=reverse)
-                result_exp = _fake_scan(RNN, h, x, dim=dim, reverse=reverse)
-                self.assertEqual(result, result_exp)
+                c_new_0 = x[0] + 1
+                c_new_1 = x[1] + 1
+                h_new = (
+                    torch.tanh(c_new_1 + x[0] @ W_hh + b_hh)
+                    + y[0] @ W_ih
+                    + y[1] @ W_ih
+                    + b_ih
+                    + x[1]
+                )
+                return (c_new_0, c_new_1), h_new
 
-                if autograd:
-                    self.check_autograd(result, result_exp, params)
+            inits = (h, h_1)
+            result = scan_fct(RNN, inits, (x, x1), dim=dim, reverse=reverse)
+            result_exp = _fake_scan(RNN, (h, h_1), (x, x1), dim=dim, reverse=reverse)
+            self.assertEqual(result, result_exp)
+
+            if autograd:
+                result_flat = pytree.tree_leaves(result)
+                result_exp_flat = pytree.tree_leaves(result_exp)
+                exp_grad_mask = [
+                    True if r.requires_grad else False for r in result_exp_flat
+                ]
+                self.check_autograd(
+                    [r for r, m in zip(result_flat, exp_grad_mask) if m],
+                    [r for r, m in zip(result_exp_flat, exp_grad_mask) if m],
+                    params,
+                )
+
+    @unittest.skipIf(not SM70OrLater, "triton")
+    @requires_cuda
+    @parametrize("reverse", [False, True])
+    @parametrize("compile_mode", ["none", "eager", "compile", "compile_dynamic_shape"])
+    @parametrize("device", [torch.device("cpu"), torch.device("cuda")])
+    @parametrize("autograd", [True])
+    def test_scan_closure_combine_fn_with_no_grad_init_carries_unequal_grad(
+        self, reverse, compile_mode, device, autograd
+    ):
+        dim = 1
+        scan_fct = compile_mode_helper(scan, compile_mode)
+        x = torch.randn(3, 10, 7, device=device, requires_grad=autograd)
+        h1 = torch.randn(3, 7, device=device, requires_grad=autograd)
+        h2 = torch.randn(3, 7, device=device, requires_grad=autograd)
+
+        with self.assertRaisesRegex(
+            RuntimeError,
+            "The init and carries need to have the same require_grad structure!.*",
+        ):
+            result = scan_fct(
+                get_scan_combine_fn("fct_c1_no_grad", False),
+                (h1, h2),
+                x,
+                dim=dim,
+                reverse=reverse,
+            )
 
     @unittest.skipIf(not SM70OrLater, "triton")
     @requires_cuda
@@ -2945,37 +3060,193 @@ def RNN(x: torch.Tensor, y: torch.Tensor):
     @parametrize("compile_mode", ["none", "eager", "compile", "compile_dynamic_shape"])
     @parametrize("device", [torch.device("cpu"), torch.device("cuda")])
     @parametrize("autograd", [False, True])
-    def test_scan_closure_combine_fn_with_no_grad(
+    def test_scan_closure_combine_fn_with_no_grad_init_carries_equal_grad(
         self, reverse, compile_mode, device, autograd
     ):
         dim = 1
         scan_fct = compile_mode_helper(scan, compile_mode)
-        x = torch.randn(3, 10, 5, device=device, requires_grad=autograd)
+        x = torch.randn(3, 10, 7, device=device, requires_grad=autograd)
+        h1 = torch.randn(3, 7, device=device, requires_grad=False)
+        h2 = torch.randn(3, 7, device=device, requires_grad=autograd)
+
+        result = scan_fct(
+            get_scan_combine_fn("fct_c1_no_grad", False),
+            (h1, h2),
+            x,
+            dim=dim,
+            reverse=reverse,
+        )
+        result_exp = _fake_scan(
+            get_scan_combine_fn("fct_c1_no_grad", False),
+            (h1, h2),
+            x,
+            dim=dim,
+            reverse=reverse,
+        )
+        self.assertEqual(result, result_exp)
+
+        if autograd:
+            # TODO: Ideally we should be able to select the results that require gradients like this
+            # [leaf for leaf in pytree.tree_leaves(result) if leaf.requires_grad == True]
+            # However, for the scan operator this does not work, as all outputs always have
+            # grad_fn=<ScanAutogradOpBackward>
+            res_req_grad_flat = pytree.tree_leaves(result)[1:]
+            res_exp_req_grad_flat = pytree.tree_leaves(result_exp)[1:]
+            self.check_autograd(res_req_grad_flat, res_exp_req_grad_flat, (x, h2))
+
+    @unittest.skipIf(not SM70OrLater, "triton")
+    @requires_cuda
+    @parametrize("reverse", [False, True])
+    @parametrize("compile_mode", ["none", "eager", "compile", "compile_dynamic_shape"])
+    @parametrize("device", [torch.device("cpu"), torch.device("cuda")])
+    @parametrize("autograd", [False, True])
+    def test_scan_closure_combine_fn_with_no_grad_for_out(
+        self, reverse, compile_mode, device, autograd
+    ):
+        dim = 1
+        scan_fct = compile_mode_helper(scan, compile_mode)
+        x = torch.randn(3, 10, 7, device=device, requires_grad=autograd)
+        h1 = torch.randn(3, 7, device=device, requires_grad=autograd)
+        h2 = torch.randn(3, 7, device=device, requires_grad=autograd)
+
+        def fct_ys_no_grad(x: torch.Tensor, y: torch.Tensor):
+            c1 = x[0] + y
+            c2 = x[1] + y
+            with torch.no_grad():
+                h_new = torch.tanh(x[0] + x[1] + y)
+            return (c1, c2), h_new
+
+        result = scan_fct(fct_ys_no_grad, (h1, h2), x, dim=dim, reverse=reverse)
+        result_exp = _fake_scan(fct_ys_no_grad, (h1, h2), x, dim=dim, reverse=reverse)
+        self.assertEqual(result, result_exp)
+
+        if autograd:
+            self.check_autograd(result[0], result_exp[0], (x, h1, h2))
+
+    @unittest.skipIf(not SM70OrLater, "triton")
+    @requires_cuda
+    @parametrize("reverse", [False, True])
+    @parametrize("compile_mode", ["none", "eager", "compile", "compile_dynamic_shape"])
+    @parametrize("device", [torch.device("cpu"), torch.device("cuda")])
+    @parametrize("autograd", [False, True])
+    def test_scan_closure_combine_fn_with_no_grad_additional_inputs_partial(
+        self, reverse, compile_mode, device, autograd
+    ):
+        dim = 1
+        scan_fct = compile_mode_helper(scan, compile_mode)
+        x = torch.randn(3, 10, 7, device=device, requires_grad=autograd)
         h = torch.randn(3, 7, device=device, requires_grad=autograd)
-        W_ih = torch.randn(5, 7, device=device, requires_grad=autograd)
+        W_ih = torch.randn(7, 7, device=device, requires_grad=autograd)
         b_ih = torch.randn(7, device=device, requires_grad=autograd)
         W_hh = torch.randn(7, 7, device=device, requires_grad=autograd)
         b_hh = torch.randn(7, device=device, requires_grad=autograd)
 
-        def RNN(x: torch.Tensor, y: torch.Tensor):
-            c_new = y @ W_ih + b_ih
+        def fct_no_grad_bhh_Whh(x: torch.Tensor, y: torch.Tensor):
+            c_new = y @ W_ih + b_ih + x
+
+            h_new = c_new + 1
             with torch.no_grad():
-                h_new = torch.tanh(c_new + x @ W_hh + b_hh)
-            return c_new, h_new
+                h_new_no_grad = torch.tanh(x @ W_hh + b_hh)
+            h_new2 = h_new + h_new_no_grad
+
+            return c_new, h_new2
+
+        result = scan_fct(fct_no_grad_bhh_Whh, h, x, dim=dim, reverse=reverse)
+        result_exp = _fake_scan(fct_no_grad_bhh_Whh, h, x, dim=dim, reverse=reverse)
+        self.assertEqual(result, result_exp)
+
+        if autograd:
+            self.check_autograd(result[1], result_exp[1], (h, x, W_ih, b_ih))
+
+    @unittest.skipIf(not SM70OrLater, "triton")
+    @requires_cuda
+    @parametrize("reverse", [False, True])
+    @parametrize("compile_mode", ["none", "eager", "compile", "compile_dynamic_shape"])
+    @parametrize("device", [torch.device("cpu"), torch.device("cuda")])
+    @parametrize("autograd", [False, True])
+    def test_scan_closure_combine_fn_with_no_grad_additional_inputs_all(
+        self, reverse, compile_mode, device, autograd
+    ):
+        dim = 1
+        scan_fct = compile_mode_helper(scan, compile_mode)
+        x = torch.randn(3, 10, 7, device=device, requires_grad=autograd)
+        h = torch.randn(3, 7, device=device, requires_grad=autograd)
+        W_ih = torch.randn(7, 7, device=device, requires_grad=autograd)
+        b_ih = torch.randn(7, device=device, requires_grad=autograd)
+        W_hh = torch.randn(7, 7, device=device, requires_grad=autograd)
+        b_hh = torch.randn(7, device=device, requires_grad=autograd)
+
+        def fct_no_grad_bih_Wih_bhh_Whh(x: torch.Tensor, y: torch.Tensor):
+            c_new = x + y
+            h_new = c_new + x
+            with torch.no_grad():
+                c_new_no_grad = y @ W_ih + b_ih
+                h_new_no_grad = torch.tanh(x @ W_hh + b_hh)
+            c_new2 = c_new + c_new_no_grad
+            h_new2 = h_new + h_new_no_grad
+            return c_new2, h_new2
+
+        result = scan_fct(fct_no_grad_bih_Wih_bhh_Whh, h, x, dim=dim, reverse=reverse)
+        result_exp = _fake_scan(
+            fct_no_grad_bih_Wih_bhh_Whh, h, x, dim=dim, reverse=reverse
+        )
+        self.assertEqual(result, result_exp)
 
         if autograd:
+            self.check_autograd(result[1], result_exp[1], (h, x))
+
+    @unittest.skipIf(not SM70OrLater, "triton")
+    @requires_cuda
+    @parametrize("reverse", [False, True])
+    @parametrize("compile_mode", ["none", "eager", "compile", "compile_dynamic_shape"])
+    @parametrize("device", [torch.device("cpu"), torch.device("cuda")])
+    @parametrize("autograd", [False, True])
+    def test_scan_closure_combine_fn_carries_ys_same_grad(
+        self, reverse, compile_mode, device, autograd
+    ):
+        dim = 1
+        scan_fct = compile_mode_helper(scan, compile_mode)
+        x = torch.randn(3, 10, 7, device=device, requires_grad=autograd)
+        h = torch.randn(3, 7, device=device, requires_grad=autograd)
+        W_ih = torch.randn(7, 7, device=device, requires_grad=autograd)
+        b_ih = torch.randn(7, device=device, requires_grad=autograd)
+        W_hh = torch.randn(7, 7, device=device, requires_grad=autograd)
+        b_hh = torch.randn(7, device=device, requires_grad=autograd)
+
+        def fct_no_grad_bih_Wih_bhh_Whh(x: torch.Tensor, y: torch.Tensor):
+            c_new = x + y
+            h_new = c_new + 1
+            with torch.no_grad():
+                c_new_no_grad = y @ W_ih + b_ih
+                h_new_no_grad = torch.tanh(x @ W_hh + b_hh)
+            c_new2 = c_new + c_new_no_grad
+            h_new2 = h_new + h_new_no_grad
+            return c_new2, h_new2
+
+        # The gradient of fct_no_grad_bih_Wih_bhh_Whh has aliased outputs
+        # which is not currently supported in inductor.
+        if compile_mode not in ["none", "eager"] and autograd:
             with self.assertRaisesRegex(
-                RuntimeError,
-                "scan currently only supports Autograd if all.*",
+                # Should be
+                # UnsupportedAliasMutationException,
+                # "Combine_fn might be aliasing its outputs!.*",
+                torch._dynamo.exc.BackendCompilerFailed,
+                "backend='inductor' raised.*",
             ):
-                result = scan_fct(RNN, h, x, dim=dim, reverse=reverse)
+                result = scan_fct(
+                    fct_no_grad_bih_Wih_bhh_Whh, h, x, dim=dim, reverse=reverse
+                )
         else:
-            result = scan_fct(RNN, h, x, dim=dim, reverse=reverse)
-            result_exp = _fake_scan(RNN, h, x, dim=dim, reverse=reverse)
+            result = scan_fct(
+                fct_no_grad_bih_Wih_bhh_Whh, h, x, dim=dim, reverse=reverse
+            )
+            result_exp = _fake_scan(
+                fct_no_grad_bih_Wih_bhh_Whh, h, x, dim=dim, reverse=reverse
+            )
             self.assertEqual(result, result_exp)
 
             if autograd:
-                self.check_autograd(result, result_exp, (x, h, W_ih, b_ih, W_hh, b_hh))
+                self.check_autograd(result[1], result_exp[1], (h, x))
 
     @unittest.skipIf(not SM70OrLater, "triton")
     @requires_cuda
@@ -3091,7 +3362,7 @@ def f(fct, init, xs):
             "Observed exception.*",
         ):
             gm = make_fx(f, tracing_mode="symbolic")(
-                get_scan_combine_fn("add", True), init, x
+                get_scan_combine_fn("add", False), init, x
             )
 
     @skipIfNoDynamoSupport
@@ -3148,7 +3419,7 @@ def f(fct, init, xs):
 
         # Correct case
         gm = make_fx(f, tracing_mode="symbolic")(
-            get_scan_combine_fn("add", False), init, x
+            get_scan_combine_fn("add", True), init, x
         )
         self.assertExpectedInline(
             gm.code.strip(),
@@ -3168,7 +3439,7 @@ def forward(self, fct_1, init_1, xs_1):
 
         # Check graph
         backend = EagerAndRecordGraphs()
-        torch.compile(f, backend=backend)(get_scan_combine_fn("add", False), init, x)
+        torch.compile(f, backend=backend)(get_scan_combine_fn("add", True), init, x)
         gm = backend.graphs[0]
 
         self.assertExpectedInline(
diff --git a/torch/_dynamo/trace_rules.py b/torch/_dynamo/trace_rules.py
index ed81b26191b4d..bd1af47063464 100644
--- a/torch/_dynamo/trace_rules.py
+++ b/torch/_dynamo/trace_rules.py
@@ -2284,7 +2284,7 @@
         "torch._guards.compile_context",
         "torch._guards.detect_fake_mode",
         "torch._guards.tracing",
-        "torch._higher_order_ops.map._has_potential_branch_input_alias",
+        "torch._higher_order_ops.map._has_potential_branch_input_output_alias",
         "torch._higher_order_ops.map._has_potential_branch_input_mutation",
         "torch._higher_order_ops.map._stack_pytree",
         "torch._higher_order_ops.map._unstack_pytree",
diff --git a/torch/_higher_order_ops/cond.py b/torch/_higher_order_ops/cond.py
index 0467e2899adc2..aa820356948c7 100644
--- a/torch/_higher_order_ops/cond.py
+++ b/torch/_higher_order_ops/cond.py
@@ -16,8 +16,8 @@
 from torch._functorch.utils import exposed_in
 from torch._guards import detect_fake_mode
 from torch._higher_order_ops.utils import (
-    _has_potential_branch_input_alias,
     _has_potential_branch_input_mutation,
+    _has_potential_branch_input_output_alias,
     _maybe_run_with_interpreter,
     _set_compilation_env,
     reenter_make_fx,
@@ -469,7 +469,7 @@ def cond_func(ctx, pred, true_fn, false_fn, inputs):
                     "One of torch.cond branch might be modifying the input!"
                 )
         for branch in [true_fn, false_fn]:
-            if _has_potential_branch_input_alias(
+            if _has_potential_branch_input_output_alias(
                 branch, unwrapped_inputs, pre_dispatch=pre_dispatch
             ):
                 raise UnsupportedAliasMutationException(
diff --git a/torch/_higher_order_ops/hints_wrap.py b/torch/_higher_order_ops/hints_wrap.py
index c211d40561463..bd6faa9aedf28 100644
--- a/torch/_higher_order_ops/hints_wrap.py
+++ b/torch/_higher_order_ops/hints_wrap.py
@@ -3,8 +3,8 @@
 import torch.utils._pytree as pytree
 from torch._C import DispatchKey
 from torch._higher_order_ops.utils import (
-    _has_potential_branch_input_alias,
     _has_potential_branch_input_mutation,
+    _has_potential_branch_input_output_alias,
     autograd_not_implemented,
     reenter_make_fx,
     unique_graph_id,
@@ -103,7 +103,7 @@ def hints_wrapper_functionalize(ctx, body_fn, args, kwargs, hints):
             raise UnsupportedAliasMutationException(
                 "body_fn of hints_wrapper might be modifying the input!"
             )
-        if _has_potential_branch_input_alias(
+        if _has_potential_branch_input_output_alias(
             functional_body_fn, unwrapped_args, pre_dispatch=pre_dispatch
         ):
             raise UnsupportedAliasMutationException(
diff --git a/torch/_higher_order_ops/map.py b/torch/_higher_order_ops/map.py
index d57d68d5e473f..244568395362f 100644
--- a/torch/_higher_order_ops/map.py
+++ b/torch/_higher_order_ops/map.py
@@ -5,8 +5,8 @@
 from torch._dispatch.python import suspend_functionalization
 from torch._functorch.aot_autograd import AOTConfig, create_joint
 from torch._higher_order_ops.utils import (
-    _has_potential_branch_input_alias,
     _has_potential_branch_input_mutation,
+    _has_potential_branch_input_output_alias,
     _maybe_run_with_interpreter,
     reenter_make_fx,
     UnsupportedAliasMutationException,
@@ -255,7 +255,7 @@ def map_functionalize(ctx, f, xs, pos_args):
         ):
             raise UnsupportedAliasMutationException("torch.map is mutating the input!")
 
-        if _has_potential_branch_input_alias(
+        if _has_potential_branch_input_output_alias(
             f, example_inputs, pre_dispatch=pre_dispatch
         ):
             raise UnsupportedAliasMutationException("torch.map is aliasing the input!")
diff --git a/torch/_higher_order_ops/scan.py b/torch/_higher_order_ops/scan.py
index 972428d7273f1..2df8b22199cec 100644
--- a/torch/_higher_order_ops/scan.py
+++ b/torch/_higher_order_ops/scan.py
@@ -11,8 +11,8 @@
 from torch._C import DispatchKey
 from torch._dispatch.python import suspend_functionalization
 from torch._higher_order_ops.utils import (
-    _has_potential_branch_input_alias,
     _has_potential_branch_input_mutation,
+    _has_potential_branch_input_output_alias,
     _set_compilation_env,
     reenter_make_fx,
     unique_graph_id,
@@ -66,6 +66,14 @@ def shift_source_dim_to_target_dim(t, from_dim: int, to_dim: int):
     return t.permute(*dims)
 
 
+def get_gradient_mask(tensor_list):
+    return [True if v is not None and v.requires_grad else False for v in tensor_list]
+
+
+def mask_gradient(grads, mask):
+    return [g for g, m in zip(grads, mask) if m]
+
+
 # Internal functions for scan.py
 def wrap_combine_fn_flat(
     *args, combine_fn, spec_init, spec_xs, num_init_leaves, num_inp_leaves
@@ -108,18 +116,26 @@ def wrapper_fwd_combine_fn(*args):
                 new_carry, y = _extract_carry_and_out(combine_fn(*args), num_init)
                 return [*new_carry, *[n_c.clone().detach() for n_c in new_carry], *y]
 
+            # The forward graph needs to be constructed from a different combine_fn than the joint_graph
+            fw_graph = _maybe_reenter_make_fx(wrapper_fwd_combine_fn)(
+                *fw_init, *fw_xs, *fw_additional_inputs
+            )
+
+            # Get gradient masks of inits, xs and additional_inputs.
+            # These masks are used during the backward path to
+            # return gradients for inits with requires_grad=True and None for the others
+            init_mask = get_gradient_mask(init)
+            xs_mask = get_gradient_mask(xs)
+            additional_inputs_mask = get_gradient_mask(additional_inputs)
+
             carry, outs = _extract_carry_and_out(
-                wrapper_fwd_combine_fn(*fw_init, *fw_xs, *fw_additional_inputs),
+                wrapper_fwd_combine_fn(
+                    *fw_init,
+                    *fw_xs,
+                    *fw_additional_inputs,
+                ),
                 num_init,
             )
-            # TODO: Support this in the future
-            if pytree.tree_any(
-                lambda t: not t.requires_grad,  # type: ignore[union-attr]
-                combine_fn(*fw_init, *fw_xs, *fw_additional_inputs),
-            ):
-                raise RuntimeError(
-                    "scan currently only supports Autograd if all init, xs and lifted parameters require gradients."
-                )
 
             fw_carry, fw_outputs = [pytree.tree_map(_from_fun, c) for c in carry], [
                 pytree.tree_map(_from_fun, o) for o in outs
@@ -140,31 +156,87 @@ def wrapper_fwd_combine_fn(*args):
                     f"Got types {[type(out) for out in fw_outputs]}."
                 )
 
-            # The forward graph needs to be constructed from a different combine_fn than the joint_graph
-            fw_graph = _maybe_reenter_make_fx(wrapper_fwd_combine_fn)(
-                *fw_init, *fw_xs, *fw_additional_inputs
-            )
-
+            # The joint graph is constructed with the requires_grad forced to True for the init, xs and additional_inputs
+            # This is necessary because during the backward scan, we need the g_init for the other gradients, even if
+            # we don't directly need g_init
             _, joint_graph = create_fw_bw_graph(
                 combine_fn,
                 False,
-                (*fw_init, *fw_xs, *fw_additional_inputs),
+                (
+                    *fw_init,
+                    *fw_xs,
+                    *fw_additional_inputs,
+                ),
                 (*fw_carry, *fw_outputs[num_init:]),
             )
 
+            g_c, g_xs = _extract_carry_and_out(
+                joint_graph(
+                    *fw_carry,
+                    *fw_outputs[num_init:],
+                    *fw_init,
+                    *fw_xs,
+                    *fw_additional_inputs,
+                ),
+                num_init,
+            )
+
+            # Check whether the init and the carries have the same requires_grad flags
+            # Scan enforces that the levaes of inits and of the carries require the same gradients
+            if any(
+                cm is not im or g_cm is not im
+                for g_cm, cm, im in zip(
+                    get_gradient_mask(g_c), get_gradient_mask(carry), init_mask
+                )
+            ):
+                raise RuntimeError(
+                    "The init and carries need to have the same require_grad structure! \
+                    E.g., check for `with torch.no_gard` statements in the combine_fn."
+                )
+
+            # The gradient masks are combined with the initial ones.
+            # The reason is that the combine_fn might contain ``with torch.no_grad()`` statements
+            # Thus, even if the gradients of xs or additional_inputs should be tracked,
+            # The ``torch.no_grad()`` statements may break the gradient tracking
+            xs_mask = xs_mask and get_gradient_mask(
+                g_xs[: len(g_xs) - num_additional_inputs]
+            )
+            additional_inputs_mask = additional_inputs_mask and get_gradient_mask(
+                g_xs[len(g_xs) - num_additional_inputs :]
+            )
+
             def wrapper_bwd_combine_fn(*args):
                 carried_g_additional_input = args[:num_additional_inputs]
+
                 g_c, g_xs = _extract_carry_and_out(
                     joint_graph(*args[num_additional_inputs:]), num_init
                 )
                 current_g_additional_inputs = g_xs[len(g_xs) - num_additional_inputs :]
+
                 new_g_additional_inputs = [
-                    carr_g + curr_g
-                    for carr_g, curr_g in zip(
-                        carried_g_additional_input, current_g_additional_inputs
+                    # The clone().detach() is required to avoid aliasing inputs
+                    carr_g + curr_g if add_inp_m else carr_g.clone().detach()
+                    for add_inp_m, carr_g, curr_g in zip(
+                        additional_inputs_mask,
+                        carried_g_additional_input,
+                        current_g_additional_inputs,
                     )
                 ]
                 g_xs = g_xs[: len(g_xs) - num_additional_inputs]
+
+                # We need to mask the g_xs so that no None values are returned
+                # The reason being that in the backward implementation, we store
+                # The gradients in a matrix and thus, None values are problematic
+                g_xs = mask_gradient(g_xs, xs_mask)
+                g_c = [
+                    g if g_m else torch.zeros_like(gi)
+                    for g, g_m, gi in zip(
+                        g_c,
+                        init_mask,
+                        args[num_additional_inputs : num_additional_inputs + num_init],
+                    )
+                ]
+
                 return [*new_g_additional_inputs, *g_c, *g_xs]
 
         new_joint_graph = _maybe_reenter_make_fx(wrapper_bwd_combine_fn)(
@@ -175,7 +247,7 @@ def wrapper_bwd_combine_fn(*args):
             *fw_xs,
             *fw_additional_inputs,
         )
-        return fw_graph, new_joint_graph
+        return fw_graph, new_joint_graph, init_mask, xs_mask, additional_inputs_mask
 
 
 def scan(
@@ -269,8 +341,14 @@ def add(x: torch.Tensor, y: torch.Tensor):
         raise RuntimeError("All init leaves must be a Tensor")
     if any(not isinstance(x, torch.Tensor) for x in leaves_xs):
         raise RuntimeError("All xs leaves must be a Tensor")
+    if any(x.ndim < dim for x in leaves_xs):
+        raise RuntimeError(
+            "All xs leaves must at least have 'dim' number of dimensions and scan dimension > 0"
+        )
     if any(x.shape[dim] == 0 for x in leaves_xs):
-        raise RuntimeError("All xs leaves must have a scan dimension > 0")
+        raise RuntimeError(
+            "All xs leaves must at least have 'dim' number of dimensions and scan dimension > 0"
+        )
 
     if len(leaves_xs) > 0:
         shape = leaves_xs[0].shape
@@ -365,19 +443,18 @@ def _scan(init, xs):
         # idxs -> Index matrix for scatter_
         # out: (num_elems, M, N, ...)
         # idx: (1, M, N)
-        outs, idxs = zip(
-            *[
-                [
-                    torch.zeros(
-                        [num_elems] + list(e.size()),
-                        dtype=e.dtype,
-                        device=e.device,
-                    ),
-                    torch.ones_like(e, dtype=torch.int64).unsqueeze(0),
-                ]
-                for i, e in enumerate(dummy_out)
-            ]
-        )
+        outs = [
+            torch.zeros(
+                [num_elems] + list(e.size()),
+                dtype=e.dtype,
+                device=e.device,
+            )
+            for i, e in enumerate(dummy_out)
+        ]
+        idxs = [
+            torch.ones_like(e, dtype=torch.int64).unsqueeze(0)
+            for i, e in enumerate(dummy_out)
+        ]
 
         def store_out_in_outs(out, ind):
             # Store the intermediate out in the outs matrix
@@ -498,6 +575,9 @@ def forward(
         reverse,
         num_leaves_init,
         num_leaves_xs,
+        carry_mask,
+        xs_mask,
+        additional_inputs_mask,
         *flat_args,
     ):
         ctx._joint_graph = joint_graph
@@ -508,7 +588,11 @@ def forward(
         init, xs, additional_inputs = ScanAutogradOp.extract_init_xs_additional_inputs(
             list(flat_args), num_leaves_init, num_leaves_xs
         )
-        ctx._num_leaves_additional_inputs = len(additional_inputs)
+        ctx._num_additional_inputs = len(additional_inputs)
+
+        ctx._carry_mask = carry_mask
+        ctx._xs_mask = xs_mask
+        ctx._additional_inputs_mask = additional_inputs_mask
 
         with torch._C._AutoDispatchBelowAutograd():
             carry, carries_outs = _extract_carry_and_out(
@@ -603,14 +687,39 @@ def prepare_initial_gradients(
             g_additional_inputs = [torch.zeros_like(ai) for ai in additional_inputs]
             return g_c_T, g_ys, g_additional_inputs
 
+        def mask_grads_with_None(real_grads, mask):
+            g_list = []
+            for m, g in zip(mask, real_grads):
+                if m:
+                    g_list.append(g)
+                else:
+                    g_list.append(None)
+            return g_list
+
+        def expand_grads_with_None(real_grads, mask):
+            g_list = []
+            real_grads_cnt = 0
+            for m in mask:
+                if m:
+                    g_list.append(real_grads[real_grads_cnt])
+                    real_grads_cnt += 1
+                else:
+                    g_list.append(None)
+            return g_list
+
         joint_graph = ctx._joint_graph
         dim = ctx._dim
         reverse = ctx._reverse
         num_leaves_init = ctx._num_leaves_init
         num_leaves_xs = ctx._num_leaves_xs
-        num_leaves_additional_inputs = ctx._num_leaves_additional_inputs
         num_leaves_ys = ctx._num_leaves_ys
 
+        carry_mask = ctx._carry_mask
+        xs_mask = ctx._xs_mask
+        num_xs_mask = sum(xs_mask)
+        additional_inputs_mask = ctx._additional_inputs_mask
+        num_additional_inputs = ctx._num_additional_inputs
+
         # The results from the forward scan are always stacked on dim 0
         # The gradients though need to be provided with the correct scan dimension dim
         # Therefore, the inputs to the backward scan are all on dim 0, and the scan is performed on dim 0
@@ -640,10 +749,6 @@ def prepare_initial_gradients(
                 num_leaves_ys,
                 bwd_scan_dim,
             )
-
-            # Prepare the inputs for the backward scan.
-            # This involves flipping the input xs if needed as well as
-            # Prepending the init of the forward scan to the carries
             xs, carries = prepare_xs_carries_for_bwd(
                 xs, init, carries, bwd_scan_dim, reverse
             )
@@ -657,15 +762,19 @@ def prepare_initial_gradients(
                 True,
                 additional_inputs,
             )
-            new_g_additional_inputs = g_outs[:num_leaves_additional_inputs]
+            new_g_additional_inputs = g_outs[:num_additional_inputs]
             g_init = g_outs[
-                num_leaves_additional_inputs : num_leaves_additional_inputs
-                + num_leaves_init
+                num_additional_inputs : num_additional_inputs + num_leaves_init
             ]
-            g_xs = g_outs[-num_leaves_xs:]
+            g_xs = g_outs[len(g_outs) - num_xs_mask :]
             g_xs = prepare_final_gradients_xs(g_xs, dim, reverse)
 
-        return *[None] * 6, *g_init, *g_xs, *new_g_additional_inputs
+        new_g_additional_inputs = mask_grads_with_None(
+            new_g_additional_inputs, additional_inputs_mask
+        )
+        g_xs = expand_grads_with_None(g_xs, xs_mask)
+        g_init = mask_grads_with_None(g_init, carry_mask)
+        return *[None] * 9, *g_init, *g_xs, *new_g_additional_inputs
 
 
 @scan_op.py_impl(DispatchKey.Autograd)
@@ -681,15 +790,6 @@ def scan_autograd(combine_fn, init, xs, dim, reverse, additional_inputs):
         with torch._C._AutoDispatchBelowAutograd():
             return scan_op(combine_fn, init, xs, dim, reverse, additional_inputs)
 
-    # TODO: Support this in the future
-    if pytree.tree_any(
-        lambda t: not t.requires_grad,  # type: ignore[union-attr]
-        (init, xs, additional_inputs),
-    ):
-        raise RuntimeError(
-            "scan currently only supports Autograd if all init, xs and lifted parameters require gradients."
-        )
-
     # TODO: The create_fw_bw is always invoked twice:
     # Once in the forward path and
     # once in the backward path, where it should only be invoked for the grad grad case.
@@ -709,6 +809,9 @@ def scan_autograd(combine_fn, init, xs, dim, reverse, additional_inputs):
     (
         fw_graph,
         joint_graph,
+        carry_mask,
+        xs_mask,
+        additional_inputs_mask,
     ) = create_fw_bw_graph_combinefn(combine_fn, init, xs, dim, additional_inputs)
 
     flat_out = ScanAutogradOp.apply(
@@ -718,6 +821,9 @@ def scan_autograd(combine_fn, init, xs, dim, reverse, additional_inputs):
         reverse,
         num_leaves_init,
         num_leaves_xs,
+        carry_mask,
+        xs_mask,
+        additional_inputs_mask,
         *(init + xs + additional_inputs),
     )
     return *flat_out[:num_leaves_init], *flat_out[num_leaves_init:]
@@ -771,13 +877,13 @@ def scan_functionalize(ctx, combine_fn, init, xs, dim, reverse, additional_input
             functional_combine_fn, sample_inputs, pre_dispatch=pre_dispatch
         ):
             raise UnsupportedAliasMutationException(
-                "Combine_fn might be modifying the input!"
+                "Combine_fn might be modifying the input! Please also check the gradient of combine_fn for modifying the input"
             )
-        if _has_potential_branch_input_alias(
+        if _has_potential_branch_input_output_alias(
             functional_combine_fn, sample_inputs, pre_dispatch=pre_dispatch
         ):
             raise UnsupportedAliasMutationException(
-                "Combine_fn might be aliasing the input!"
+                "Combine_fn might be aliasing the input! Please also check the gradient of combine_fn for modifying the input"
             )
         ret = scan_op(
             functional_combine_fn,
diff --git a/torch/_higher_order_ops/utils.py b/torch/_higher_order_ops/utils.py
index 139e9a160cbe2..0479a5b78eef3 100644
--- a/torch/_higher_order_ops/utils.py
+++ b/torch/_higher_order_ops/utils.py
@@ -114,12 +114,7 @@ def _set_compilation_env():
         torch.fx._symbolic_trace._is_fx_tracing_flag = _old_is_tracing
 
 
-def _has_potential_branch_input_mutation(branch, inputs, pre_dispatch=False):
-    """
-    Dispatch-trace the branch with inputs and check if
-    producing graph has mutable op on the input. This is
-    bit restrictive as the branch must be traceable.
-    """
+def _create_gm_from_branch(branch, inputs, pre_dispatch):
     try:
         gm = make_fx(branch, pre_dispatch=pre_dispatch)(*inputs)
     except UnsupportedAliasMutationException:
@@ -129,6 +124,21 @@ def _has_potential_branch_input_mutation(branch, inputs, pre_dispatch=False):
     except Exception as e:
         raise e
 
+    return gm
+
+
+def _has_potential_branch_input_mutation(branch, inputs, pre_dispatch=False):
+    """
+    Dispatch-trace the branch with inputs and check if
+    producing graph has mutable op on the input. This is
+    bit restrictive as the branch must be traceable.
+    """
+    gm = _create_gm_from_branch(branch, inputs, pre_dispatch)
+    # this can happen when nested cond_op is
+    # functionalized
+    if isinstance(gm, bool):
+        return gm
+
     def _detect_input_mutation(gm):
         input_nodes = set()
         for node in gm.graph.nodes:
@@ -154,20 +164,17 @@ def _detect_input_mutation(gm):
     return _detect_input_mutation(gm)
 
 
-def _has_potential_branch_input_alias(branch, inputs, pre_dispatch=False):
+def _has_potential_branch_input_output_alias(branch, inputs, pre_dispatch=False):
     """
     Dispatch-trace the branch with inputs and check if
     producing graph has output aliasing the branch input. This is
     bit restrictive as the branch must be traceable.
     """
-    try:
-        gm = make_fx(branch, pre_dispatch=pre_dispatch)(*inputs)
-    except UnsupportedAliasMutationException:
-        # this can happen when nested cond_op is
-        # functionalized
-        return True
-    except Exception as e:
-        raise e
+    gm = _create_gm_from_branch(branch, inputs, pre_dispatch)
+    # this can happen when nested cond_op is
+    # functionalized
+    if isinstance(gm, bool):
+        return gm
 
     def _detect_input_alias(gm):
         input_storages = set()
@@ -178,11 +185,16 @@ def _detect_input_alias(gm):
             if node.op == "placeholder" and "val" in node.meta:
                 input_storages.add(StorageWeakRef(node.meta["val"]._typed_storage()))
             if node.op == "output":
+                output_storages = set()
 
                 def check_alias(out):
                     if out is not None and "val" in out.meta:
                         out_storage = StorageWeakRef(out.meta["val"]._typed_storage())
-                        return out_storage in input_storages
+                        alias = (out_storage in input_storages) or (
+                            out_storage in output_storages
+                        )
+                        output_storages.add(out_storage)
+                        return alias
                     return False
 
                 if any(pytree.tree_leaves(pytree.tree_map(check_alias, node.args))):
@@ -213,7 +225,7 @@ def unique_graph_id(proxy_mode, prefix):
     return i, next_name
 
 
-def _from_fun(t):
+def _from_fun(t, force_requires_grad=False):
     from torch._functorch.aot_autograd import from_fun
     from torch._subclasses.functional_tensor import FunctionalTensor
 
@@ -223,7 +235,7 @@ def _from_fun(t):
                 t.size(),
                 t.stride(),
                 dtype=t.dtype,
-                requires_grad=t.requires_grad,
+                requires_grad=t.requires_grad if not force_requires_grad else True,
                 device=t.device,
             )
         else:
diff --git a/torch/_higher_order_ops/while_loop.py b/torch/_higher_order_ops/while_loop.py
index f14321842f40b..04cfc8c152a75 100644
--- a/torch/_higher_order_ops/while_loop.py
+++ b/torch/_higher_order_ops/while_loop.py
@@ -5,8 +5,8 @@
 import torch.utils._pytree as pytree
 from torch._C import DispatchKey
 from torch._higher_order_ops.utils import (
-    _has_potential_branch_input_alias,
     _has_potential_branch_input_mutation,
+    _has_potential_branch_input_output_alias,
     _maybe_run_with_interpreter,
     _set_compilation_env,
     autograd_not_implemented,
@@ -266,7 +266,7 @@ def while_loop_func(ctx, cond_fn, body_fn, carried_inputs, additional_inputs):
                     f"torch.while_loop's {fn_name} might be modifying the input!"
                 )
 
-            if _has_potential_branch_input_alias(
+            if _has_potential_branch_input_output_alias(
                 fn, unwrapped_inputs, pre_dispatch=pre_dispatch
             ):
                 raise UnsupportedAliasMutationException(