Changes intermediate

Author: Dhruvanshu-Joshi

pytensor/tensor/rewriting/subtensor.py

@@ -406,26 +406,24 @@ def local_replace_slice(fgraph, node):
         if not isinstance(s, slice):
             continue
 
-        idx_change = False
         start = s.start
         stop = s.stop
         step = s.step
         if extract_constant(start, only_process_constants=True) == 0:
-            idx_change = True
+            idx_flag = True
             start = None
 
         if extract_constant(stop, only_process_constants=True) == x.type.shape[dim]:
-            idx_change = True
+            idx_flag = True
             stop = None
 
         if extract_constant(step, only_process_constants=True) == 1:
-            idx_change = True
+            idx_flag = True
             step = None
 
-        if idx_change:
-            idx_flag = True
-            new_idxs[dim] = slice(start, stop, step)
+        new_idxs[dim] = slice(start, stop, step)
 
+    assert node.outputs[0].type == x[tuple(new_idxs)].type
     if idx_flag is True:
         return [x[tuple(new_idxs)]]
     else:

tests/tensor/rewriting/test_subtensor.py

@@ -2411,6 +2411,7 @@ def test_slice_canonicalize(fstop, lstop, lstep):
     f = pytensor.function([x], y)
 
     test_y = f(np.random.normal(size=(3, 5, 6, 9)))
+    # test_y = f.maker.fgraph.toposort()
 
     y1 = x[None:None:None, None:None:None, None:7:None, None:None:-1]
 
@@ -2420,6 +2421,7 @@ def test_slice_canonicalize(fstop, lstop, lstep):
     f1 = pytensor.function([x], y1)
 
     expected_y = f1(np.random.normal(size=(3, 5, 6, 9)))
+    # expected_y = f1.maker.fgraph.toposort()
 
     # TODO: test using equal_computations
     assert test_y.all() == expected_y.all()