Do not try to save initial values buffer size in Scan

This will always require a roll at the end, for a minimal gain

Author: ricardoV94

pytensor/scan/rewriting.py

@@ -1430,9 +1430,18 @@ def scan_save_mem(fgraph, node):
                     store_steps[i] = 0
                     break
 
-                if isinstance(this_slice[0], slice) and this_slice[0].start is None:
-                    store_steps[i] = 0
-                    break
+                if isinstance(this_slice[0], slice):
+                    start = this_slice[0].start
+                    if isinstance(start, Constant):
+                        start = start.data
+                    # Don't do anything if the subtensor is starting from the beginning of the buffer
+                    # Or just skipping the initial values (default output returned to the user).
+                    # Trimming the initial values would require a roll to align the buffer once scan is done
+                    # As it always starts writing at position [0+max(taps)], and ends up at position [:max(taps)]
+                    # It's cheaper to just keep the initial values in the buffer and slice them away (default output)
+                    if start in (0, None, init_l[i]):
+                        store_steps[i] = 0
+                        break
 
                 # Special case for recurrent outputs where only the last result
                 # is requested. This is needed for this rewrite to apply to

tests/link/numba/test_scan.py

@@ -476,7 +476,7 @@ def buffer_tester(self, n_steps, op_size, buffer_size, benchmark=None):
             expected_buffer_size = 3
         elif buffer_size == "whole":
             xs_kept = xs  # What users think is the whole buffer
-            expected_buffer_size = n_steps - 1
+            expected_buffer_size = n_steps
         elif buffer_size == "whole+init":
             xs_kept = xs.owner.inputs[0]  # Whole buffer actually used by Scan
             expected_buffer_size = n_steps