Fix scan_save_mem uninitialized memory read when n_steps=0 (fixes #1878)

pytensor/scan/rewriting.py

@@ -1467,12 +1467,9 @@ def scan_save_mem_rewrite(fgraph, node, backend_supports_output_pre_allocation:
 
         # FIXME: This is not correct. Scan with 0 steps seems to be supported
         # Make sure the ScanSaveMem optimization never makes the new
-        # number of steps to be 0 (this could happen, for instance, if
-        # the optimization detects that the outputs of the Scan go through
-        # subtensor nodes that end up taking no elements) because Scan with
-        # 0 iterations are not supported. Make sure the new number of steps
-        # is at least 1.
-        nw_steps = select_max(nw_steps, 1)
+        # number of steps to be 0 ... because Scan with 0 iterations are not supported.
+        # Make sure the new number of steps is at least 1, UNLESS explicitly requested as 0.
+        nw_steps = pt.switch(pt.eq(node.inputs[0], 0), 0, select_max(nw_steps, 1))
 
     # 2.4 Loop over the clients again now looking just to see how many
     # intermediate steps to store
@@ -1764,10 +1761,19 @@ def scan_save_mem_rewrite(fgraph, node, backend_supports_output_pre_allocation:
                             isinstance(old_slices[0], ScalarConstant)
                             and old_slices[0].value == -1
                         ):
-                            position = old_slices[0]
+                            # Bypass the padded uninitialized slot when n_steps=0 by fetching the last initialized tap
+                            position = pt.switch(
+                                pt.eq(node.inputs[0], 0), init_l[pos] - 1, old_slices[0]
+                            )
                         else:
-                            position = (
-                                cnf_slice[0] - nw_steps - init_l[pos] + store_steps[pos]
+                            # Bypass the buffer shift when n_steps=0, mapping directly to the canonical tap index
+                            position = pt.switch(
+                                pt.eq(node.inputs[0], 0),
+                                cnf_slice[0],
+                                cnf_slice[0]
+                                - nw_steps
+                                - init_l[pos]
+                                + store_steps[pos],
                             )
 
                         nw_slice = (sanitize(position), *old_slices[1:])

tests/scan/test_rewriting.py

@@ -1790,6 +1790,26 @@ def test_broadcasted_init(self, keep_beginning, val_ndim):
         )
         assert buffer_size_fn(val_test) == 52 if keep_beginning else 50
 
+    def test_zero_steps_uninitialized_memory(self):
+        # Regression test for PyMC #7380
+        n = pt.tensor("n", shape=(), dtype=int)
+        init_state = pt.tensor("init_state", shape=(3,))
+
+        final_state = scan(
+            fn=lambda xtm1: xtm1 * 2,
+            outputs_info=[init_state],
+            n_steps=n,
+            return_updates=False,
+        )
+        res = final_state.owner.inputs[0][-1]
+
+        fn = function([init_state, n], res, mode=self.mode, on_unused_input="ignore")
+
+        np.testing.assert_allclose(fn(init_state=np.ones(3), n=0), np.ones(3))
+        np.testing.assert_allclose(fn(init_state=np.ones(3), n=1), np.full(3, 2.0))
+        np.testing.assert_allclose(fn(init_state=np.ones(3), n=2), np.full(3, 4.0))
+        np.testing.assert_allclose(fn(init_state=np.ones(3), n=3), np.full(3, 8.0))
+
 
 def test_inner_replace_dot():
     """