Cleanup Scan symbolic buffer size graph

Graph was being broken by Scalar/Tensor conversions that prevented fusion

Author: ricardoV94

Diff for: pytensor/scan/rewriting.py

@@ -69,6 +69,7 @@
     get_idx_list,
     get_slice_elements,
     set_subtensor,
+    undo_scalarization,
 )
 from pytensor.tensor.variable import TensorConstant, TensorVariable
 
@@ -1343,6 +1344,8 @@ def scan_save_mem_rewrite(fgraph, node, backend_supports_output_pre_allocation:
                     except KeyError:
                         length = out.shape[0]
                 cf_slice = get_canonical_form_slice(this_slice[0], length)
+                cf_slice = (undo_scalarization(cf_slice[0]), cf_slice[1])
+
                 slices[i] += [(cf_slice, this_slice)]  # type: ignore
 
                 if isinstance(this_slice[0], slice) and this_slice[0].stop is None:

Diff for: pytensor/tensor/rewriting/subtensor.py

@@ -84,6 +84,7 @@
     get_slice_elements,
     inc_subtensor,
     indices_from_subtensor,
+    undo_scalarization,
 )
 from pytensor.tensor.type import TensorType, integer_dtypes
 from pytensor.tensor.type_other import NoneTypeT, SliceConstant, SliceType
@@ -1136,6 +1137,7 @@ def merge_two_slices(fgraph, slice1, len1, slice2, len2):
             # We are in the more complex case when we do not actually know
             # if the first slice was in reverse or not.
             # in case it was not in reverse:
+            sl2 = undo_scalarization(sl2)
             p_val = sl1.start + sl2 * sl1.step
             # case it was in reverse we need to realize that we do not want
             # the k-th element from sl.start but the k-th element from

Diff for: pytensor/tensor/subtensor.py

@@ -35,6 +35,9 @@
     nonzero,
     scalar_from_tensor,
 )
+from pytensor.tensor.basic import (
+    constant as tensor_constant,
+)
 from pytensor.tensor.blockwise import vectorize_node_fallback
 from pytensor.tensor.elemwise import DimShuffle
 from pytensor.tensor.exceptions import AdvancedIndexingError, NotScalarConstantError
@@ -252,6 +255,23 @@ def get_idx_list(inputs, idx_list):
     return indices_from_subtensor(inputs[1:], idx_list)
 
 
+def undo_scalarization(x):
+    """Undo scalarization of a variable.
+
+    PyTensor Basic index operations use ScalarVariables for the indices/slice arguments.
+    When reason symbolically about the result of multiple indexing operations, we usually
+    want to work on TensorVariables, since rewrites work on those and not ScalarVariables.
+
+    This function undoes ScalarFromTensor operation or converts ScalarConstants to TensorConstants.
+    """
+    if isinstance(x, ScalarVariable):
+        if isinstance(x, ScalarConstant):
+            return tensor_constant(x.data, dtype=x.dtype)
+        elif x.owner is not None and isinstance(x.owner.op, ScalarFromTensor):
+            return x.owner.inputs[0]
+    return x
+
+
 @overload
 def get_canonical_form_slice(
     theslice: slice,
@@ -298,6 +318,7 @@ def get_canonical_form_slice(
 
     # Other non-slice types are the scalar indexing case
     if not isinstance(theslice, slice):
+        theslice = undo_scalarization(theslice)
         if isinstance(theslice, int | np.integer | ScalarVariable) or (
             isinstance(theslice, TensorVariable) and theslice.ndim == 0
         ):
@@ -381,6 +402,7 @@ def analyze(x):
         elif is_stop_length:
             # start:length:1
             if is_start_constant and start >= 0:
+                length = undo_scalarization(length)
                 return slice(switch(lt(start, length), start, length), length, 1), 1
             start_plus_len = start + length
             start = switch(