Cleanup Scan symbolic buffer size graph

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

Author: ricardoV94

pytensor/scan/rewriting.py

@@ -41,7 +41,7 @@
 from pytensor.graph.type import HasShape
 from pytensor.graph.utils import InconsistencyError
 from pytensor.raise_op import Assert
-from pytensor.scalar import ScalarConstant
+from pytensor.scalar import ScalarConstant, ScalarVariable
 from pytensor.scan.op import Scan, ScanInfo
 from pytensor.scan.utils import (
     ScanArgs,
@@ -54,6 +54,7 @@
 from pytensor.tensor.basic import (
     Alloc,
     AllocEmpty,
+    ScalarFromTensor,
     get_scalar_constant_value,
 )
 from pytensor.tensor.elemwise import DimShuffle, Elemwise
@@ -1300,7 +1301,7 @@ def scan_save_mem_rewrite(fgraph, node, backend_supports_output_pre_allocation:
     # or not
     flag_store = False
 
-    # 2.2 Loop over the clients
+    # 2.2 Loop over the clients to figure out how many steps we actually need to do in the Scan
     for i, out in enumerate(node.outputs[:c_outs]):
         # look at all its clients
         slices[i] = []
@@ -1343,6 +1344,14 @@ 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)
+
+                if (
+                    isinstance(cf_slice[0], ScalarVariable)
+                    and cf_slice[0].owner is not None
+                    and isinstance(cf_slice[0].owner.op, ScalarFromTensor)
+                ):
+                    cf_slice = (cf_slice[0].owner.inputs[0], cf_slice[1])
+
                 slices[i] += [(cf_slice, this_slice)]
 
                 if isinstance(this_slice[0], slice) and this_slice[0].stop is None:

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

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
@@ -266,6 +269,15 @@ def get_canonical_form_slice(
 ) -> tuple[ScalarVariable, int]: ...
 
 
+def undo_scalarization(x):
+    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
+
+
 def get_canonical_form_slice(
     theslice: slice | int | np.integer | ScalarVariable | TensorVariable,
     length: int | np.integer | ScalarVariable | TensorVariable,
@@ -301,6 +313,7 @@ def get_canonical_form_slice(
         if isinstance(theslice, int | np.integer | ScalarVariable) or (
             isinstance(theslice, TensorVariable) and theslice.ndim == 0
         ):
+            theslice = undo_scalarization(theslice)
             cano = switch(lt(theslice, 0), (theslice + length), theslice)
             return scalar_from_tensor(cano), 1
         raise ValueError(f"Slice {theslice} is not a supported slice type.")
@@ -381,6 +394,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(