Changes for rewriting redundant slices

Author: Dhruvanshu-Joshi

pytensor/tensor/rewriting/subtensor.py

@@ -389,56 +389,51 @@ def local_useless_slice(fgraph, node):
 def local_replace_slice(fgraph, node):
     """
     Rewrite Subtensor of the form:
-        1. X[0:-1:1] -> X[:-1]
-        2. X[0:-1] -> X[:-1]
-        3. X[:-1] -> X[:-1]
+        1. X[0:-1:1] -> X[None:None:None]
+        2. X[0:-1:2] -> X[None:None:2]
+        3. X[3:-1] -> X[3:None:None]
 
     """
     idxs = get_idx_list(node.inputs, node.op.idx_list)
     x = node.inputs[0]
 
     if not idxs:
-        return [node.inputs[0]]
-
-    # pytensor.dprint(node)
-    # print(node.inputs[0].type.shape)
+        return [x]
 
-    # last_slice = len(idxs)
     new_idxs = list(idxs)
+    idx_flag = False
+    for dim, s in enumerate(new_idxs):
+        if not isinstance(s, slice):
+            continue
+
+        flag = False
+        start = s.start
+        stop = s.stop
+        step = s.step
+        if start is None or extract_constant(start, only_process_constants=True) == 0:
+            flag = True
+            start = None
 
-    # flag = False
-    # index = -1
-    # call s slice
-    for dim, s in enumerate(idxs):
         if (
-            isinstance(s, slice)
-            and (s.start is None or extract_constant(s.start, only_process_constants=True) == 0)
-            and (extract_constant(s.stop, only_process_constants=True) == -1 or extract_constant(s.stop, only_process_constants=True) == node.inputs[0].type.shape[dim])
-            and (s.step is None or extract_constant(s.step, only_process_constants=True) == 1)
+            extract_constant(stop, only_process_constants=True) == -1
+            or extract_constant(stop, only_process_constants=True) == x.type.shape[dim]
         ):
-            # break
-            if index != -1:
-            new_idxs[dim] =  slice(None, None, None)
-        else:
-            # exchange with if
-            continue
-    # if nothing changewd, return None
-    # if index != -1:
-    #     new_idxs[dim] =  slice(None, None, None)
-
-        # new_subtensor = Subtensor(tuple(new_idxs))
-        # new_subtensor_inputs = get_slice_elements(
-        #     new_idxs, lambda x: isinstance(x, Variable)
-        # )
-        # out = new_subtensor(node.inputs[0], *new_subtensor_inputs)
-        # # Copy over previous output stacktrace
-        # copy_stack_trace(node.outputs, out)
-        # return [out]
-    if change:
-        x[tuple(new_idxs)]
+            flag = True
+            stop = None
+
+        if step is None or extract_constant(step, only_process_constants=True) == 1:
+            flag = True
+            step = None
+
+        if flag:
+            idx_flag = True
+            new_idxs[dim] = slice(start, stop, step)
+
+    if idx_flag is True:
+        return [x[tuple(new_idxs)]]
     else:
         # Subtensor is not needed at all
-        return None
+        return [x]
 
 
 # fast_compile to allow opt subtensor(cast{float32}(make_vector))

tests/tensor/rewriting/test_subtensor.py

@@ -10,7 +10,7 @@
 from pytensor.compile.ops import DeepCopyOp
 from pytensor.configdefaults import config
 from pytensor.graph import FunctionGraph, vectorize_graph
-from pytensor.graph.basic import Constant, Variable, ancestors
+from pytensor.graph.basic import Constant, Variable, ancestors, equal_computations
 from pytensor.graph.rewriting.basic import check_stack_trace
 from pytensor.graph.rewriting.db import RewriteDatabaseQuery
 from pytensor.graph.rewriting.utils import rewrite_graph
@@ -2405,11 +2405,17 @@ def test_local_blockwise_advanced_inc_subtensor(set_instead_of_inc):
 
 
 def test_slice_canonicalize():
-    # Dummy test. Need to and will be changed
     x = tensor("x", shape=(3, 5, None, 9))
     y = x[:-1, 0:5, 0:7, 0:-1:-1]
 
-    f = pytensor.function([x], [y])
-    # use equal_computations
-    pytensor.dprint(f)
-    assert 0
+    f = pytensor.function([x], y)
+
+    test_y = f.maker.fgraph.toposort()
+
+    y1 = x[None:None:None, None:None:None, None:7:None, None:None:-1]
+
+    f = pytensor.function([x], y1)
+
+    expected_y = f.maker.fgraph.toposort()
+
+    assert equal_computations([test_y], [expected_y])