Don't support symbolic slices in get_canonical_form_slice

Author: michaelosthege

pytensor/tensor/subtensor.py

@@ -216,7 +216,7 @@ def as_index_literal(
         )
 
     if not isinstance(idx, Variable):
-        raise NotScalarConstantError()
+        raise TypeError(f"Not an index element: {idx}")
 
     if isinstance(idx.type, NoneTypeT):
         return None
@@ -244,23 +244,26 @@ def get_idx_list(inputs, idx_list):
 
 @overload
 def get_canonical_form_slice(
-    theslice: slice | SliceConstant,
+    theslice: slice,
     length: int | np.integer | ScalarVariable | TensorVariable,
 ) -> tuple[slice, int | ScalarConstant]: ...
 
 
 @overload
 def get_canonical_form_slice(
-    theslice: ScalarVariable | TensorVariable,
+    theslice: int | np.integer | ScalarVariable | TensorVariable,
     length: int | np.integer | ScalarVariable | TensorVariable,
 ) -> tuple[ScalarVariable, int]: ...
 
 
 def get_canonical_form_slice(
-    theslice: slice | Variable,
+    theslice: slice | int | np.integer | ScalarVariable | TensorVariable,
     length: int | np.integer | ScalarVariable | TensorVariable,
 ) -> tuple[slice | ScalarVariable, int | ScalarConstant]:
-    """Convert slices to canonical form.
+    """Convert indices or slices to canonical form.
+
+    Symbolic indices or symbolic slice attributes are supported.
+    Symbolic slices are not.
 
     Given a slice [start:stop:step] transform it into a canonical form
     that respects the conventions imposed by python and numpy.
@@ -281,35 +284,16 @@ def get_canonical_form_slice(
     """
     from pytensor.tensor import ge, lt, sign, switch
 
-    # Convert the two symbolic slice types into a native slice
-    if isinstance(theslice, SliceConstant):
-        theslice = cast(slice, theslice.data)
-    elif isinstance(theslice, Variable) and isinstance(theslice.type, SliceType):
-        theslice = slice(*theslice.owner.inputs)
     # Other non-slice types are the scalar indexing case
-    elif not isinstance(theslice, slice):
-        sslice: int | np.integer | ScalarVariable | TensorVariable
-        if isinstance(theslice, ScalarVariable):
-            sslice = theslice
-        elif isinstance(theslice, TensorVariable) and theslice.ndim == 0:
-            sslice = theslice
-        else:
-            # Only non-variable types should remain.
-            try:
-                vlit = as_index_literal(theslice)
-                if vlit is None:
-                    raise ValueError("Can't create canonical slice for `None` slices.")
-                # Slice returns from as_index are not expected, because the
-                # symbolic slice type inputs were already taken care of above.
-                assert isinstance(vlit, int | np.integer)
-                sslice = vlit
-            except NotScalarConstantError:
-                raise ValueError(f"Slice {theslice} is not a supported slice type.")
-        cano = switch(lt(sslice, 0), (sslice + length), sslice)
-        return scalar_from_tensor(cano), 1
+    if not isinstance(theslice, slice):
+        if isinstance(theslice, int | np.integer | ScalarVariable) or (
+            isinstance(theslice, TensorVariable) and theslice.ndim == 0
+        ):
+            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.")
 
     # At this point we have a slice object. Possibly with symbolic inputs.
-    assert isinstance(theslice, slice)
 
     def analyze(x):
         try:

tests/tensor/test_subtensor.py

@@ -71,7 +71,6 @@
 from pytensor.tensor.type_other import (
     NoneConst,
     SliceConstant,
-    SliceType,
     as_symbolic_slice,
     make_slice,
     slicetype,
@@ -112,11 +111,19 @@ def test_as_index_literal():
 
 
 class TestGetCanonicalFormSlice:
-    def test_none_raises(self):
-        with pytest.raises(ValueError, match="Can't create"):
-            get_canonical_form_slice(NoneConst, 5)
-        with pytest.raises(ValueError, match="Can't create"):
-            get_canonical_form_slice(None, 5)
+    @pytest.mark.parametrize(
+        "idx",
+        [
+            NoneConst,
+            None,
+            as_symbolic_slice(slice(3, 7, 2)),
+            as_symbolic_slice(slice(3, int16(), 2)),
+            vector(),
+        ],
+    )
+    def test_unsupported_inputs(self, idx):
+        with pytest.raises(ValueError, match="not a supported slice"):
+            get_canonical_form_slice(idx, 5)
 
     def test_scalar_constant(self):
         a = as_scalar(0)
@@ -146,26 +153,6 @@ def test_symbolic_tensor(self):
         assert isinstance(res[0].owner.op, ptb.ScalarFromTensor)
         assert res[1] == 1
 
-    def test_constant_slice(self):
-        idx = as_symbolic_slice(slice(3, 7, 2))
-        assert isinstance(idx, SliceConstant)
-        res = get_canonical_form_slice(idx, 10)
-        assert isinstance(res[0], slice)
-        assert res[1] == 1
-
-    def test_symbolic_slice(self):
-        idx = as_symbolic_slice(slice(3, int16(), 2))
-        assert not isinstance(idx, slice)
-        assert isinstance(idx.type, SliceType)
-        res = get_canonical_form_slice(idx, 10)
-        assert isinstance(res[0], slice)
-        assert res[1] == 1
-
-    def test_unsupported_variable(self):
-        idx = vector()
-        with pytest.raises(ValueError, match="slice type"):
-            get_canonical_form_slice(idx, 4)
-
     def test_all_integer(self):
         res = get_canonical_form_slice(slice(1, 5, 2), 7)
         assert isinstance(res[0], slice)