Intermediate draft changes

Author: Dhruvanshu-Joshi

pytensor/tensor/rewriting/subtensor.py

@@ -381,49 +381,64 @@ def local_useless_slice(fgraph, node):
             return [node.inputs[0]]
 
 
-@register_infer_shape
+@register_useless
 @register_canonicalize
+@register_stabilize
 @register_specialize
 @node_rewriter([Subtensor])
 def local_replace_slice(fgraph, node):
     """
-    Remove Subtensor of the form:
+    Rewrite Subtensor of the form:
         1. X[0:-1:1] -> X[:-1]
         2. X[0:-1] -> X[:-1]
         3. X[:-1] -> X[:-1]
 
     """
     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)
+
     # last_slice = len(idxs)
+    new_idxs = list(idxs)
 
-    for s in idxs[::-1]:
+    # 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
-            and (
-                s.step is None
-                or extract_constant(s.step, only_process_constants=True) == 1
-            )
+            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)
         ):
-            # This does not work.
-            # I get the error that
-            # ```
-            # s.start = None
-            # AttributeError: readonly attribute
-            # ```
-            s.start = None
-            s.stop = -1
-            s.step = None
+            # break
+            if index != -1:
+            new_idxs[dim] =  slice(None, None, None)
         else:
-            break
+            # 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)]
+    else:
+        # Subtensor is not needed at all
+        return None
 
 
 # fast_compile to allow opt subtensor(cast{float32}(make_vector))

tests/tensor/rewriting/test_subtensor.py

@@ -2406,11 +2406,10 @@ def test_local_blockwise_advanced_inc_subtensor(set_instead_of_inc):
 
 def test_slice_canonicalize():
     # Dummy test. Need to and will be changed
-    x = pt.vector("x")
-    y1 = x[:-1]
-    y2 = x[0:-1]
-    y1 = x[0:-1:1]
+    x = tensor("x", shape=(3, 5, None, 9))
+    y = x[:-1, 0:5, 0:7, 0:-1:-1]
 
-    f = pytensor.function([x], [y1, y2, y3])
+    f = pytensor.function([x], [y])
+    # use equal_computations
     pytensor.dprint(f)
-    # assert 0
+    assert 0