Further type refactoring of slice helper functions

Author: michaelosthege

pytensor/tensor/subtensor.py

@@ -27,7 +27,12 @@
     as_tensor_variable,
     get_vector_length,
 )
-from pytensor.tensor.basic import alloc, get_underlying_scalar_constant_value, nonzero
+from pytensor.tensor.basic import (
+    alloc,
+    get_underlying_scalar_constant_value,
+    nonzero,
+    scalar_from_tensor,
+)
 from pytensor.tensor.blockwise import vectorize_node_fallback
 from pytensor.tensor.elemwise import DimShuffle
 from pytensor.tensor.exceptions import AdvancedIndexingError, NotScalarConstantError
@@ -52,8 +57,14 @@
     wscalar,
     zscalar,
 )
-from pytensor.tensor.type_other import NoneConst, NoneTypeT, SliceType, make_slice
-from pytensor.tensor.variable import TensorVariable
+from pytensor.tensor.type_other import (
+    NoneConst,
+    NoneTypeT,
+    SliceConstant,
+    SliceType,
+    make_slice,
+)
+from pytensor.tensor.variable import TensorConstant, TensorVariable
 
 
 _logger = logging.getLogger("pytensor.tensor.subtensor")
@@ -165,19 +176,26 @@ def as_index_literal(idx: None) -> None: ...
 
 
 @overload
-def as_index_literal(idx: slice) -> slice: ...
+def as_index_literal(idx: slice | SliceConstant) -> slice: ...
 
 
 @overload
-def as_index_literal(idx: Constant) -> int | np.integer | None: ...
+def as_index_literal(idx: ScalarConstant | TensorConstant) -> int | np.integer: ...
 
 
 @overload
-def as_index_literal(idx: Variable) -> int | np.integer | slice | None: ...
+def as_index_literal(idx: Variable): ...
 
 
 def as_index_literal(
-    idx: int | np.integer | Variable | slice | None,
+    idx: None
+    | int
+    | np.integer
+    | slice
+    | SliceConstant
+    | ScalarConstant
+    | TensorConstant
+    | Variable,
 ) -> int | np.integer | slice | None:
     """Convert a symbolic index element to its Python equivalent.
 
@@ -187,40 +205,61 @@ def as_index_literal(
     ------
     NotScalarConstantError
     """
-    if isinstance(idx, int | np.integer):
+    if idx is None or isinstance(idx, int | np.integer):
         return idx
 
-    # NOTE: np.newaxis is None
-    idxtype = getattr(idx, "type", None)
-    if idx is None or isinstance(idxtype, NoneTypeT):
-        return None
-
-    if isinstance(idx, Constant):
-        data = idx.data
-        if isinstance(data, np.ndarray):
-            return cast(int, data.item())
-        return cast(int | None, data)
-
-    if isinstance(idx, Variable) and isinstance(idxtype, SliceType):
-        idx = slice(*idx.owner.inputs)
-
     if isinstance(idx, slice):
         return slice(
             as_index_literal(idx.start),
             as_index_literal(idx.stop),
             as_index_literal(idx.step),
         )
 
+    if not isinstance(idx, Variable):
+        raise NotScalarConstantError()
+
+    if isinstance(idx.type, NoneTypeT):
+        return None
+
+    if isinstance(idx, ScalarConstant):
+        return cast(int | None, idx.data)
+
+    if isinstance(idx, TensorConstant):
+        return cast(int, idx.data.item())
+
+    if isinstance(idx, SliceConstant):
+        return cast(slice, idx.data)
+
+    if isinstance(idx.type, SliceType):
+        assert idx.owner is not None
+        return slice(*map(as_index_literal, idx.owner.inputs))
+
+    # Other kinds of variables are not supported
     raise NotScalarConstantError()
 
 
 def get_idx_list(inputs, idx_list):
     return indices_from_subtensor(inputs[1:], idx_list)
 
 
+@overload
+def get_canonical_form_slice(
+    theslice: slice | SliceConstant,
+    length: int | np.integer | ScalarVariable | TensorVariable,
+) -> tuple[slice, int | ScalarConstant]: ...
+
+
+@overload
+def get_canonical_form_slice(
+    theslice: ScalarVariable | TensorVariable,
+    length: int | np.integer | ScalarVariable | TensorVariable,
+) -> tuple[slice | ScalarVariable, int]: ...
+
+
 def get_canonical_form_slice(
-    theslice: slice | Variable, length: ScalarVariable
-) -> tuple[slice | Variable, int | ScalarConstant]:
+    theslice: slice | Variable,
+    length: int | np.integer | ScalarVariable | TensorVariable,
+) -> tuple[slice | ScalarVariable, int | ScalarConstant]:
     """Convert slices to canonical form.
 
     Given a slice [start:stop:step] transform it into a canonical form
@@ -230,36 +269,46 @@ def get_canonical_form_slice(
     in which 0 <= start <= stop <= length and step > 0, and a flag which says
     if the resulting set of numbers needs to be reversed or not.
 
+    Given a scalar index `idx` that may or not be negative, convert it to
+    a certainly positive form `idx if idx >= 0 else length + idx`.
+
     Returns
     -------
     slc
-        Slice or scalar variable.
+        Canonical form slice or scalar variable.
     direction
-        Direction to iterate the resulting elements in. (-1 or 1).
+        Direction to iterate the resulting elements in. (-1 or 1). May be symbolic.
     """
     from pytensor.tensor import ge, lt, sign, switch
 
-    if not isinstance(theslice, slice):
-        # Try to extract a scalar slice and return it already.
-        sslice: int | np.integer | ScalarVariable | None = None
+    # 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:
-                    return slice(0, length, 1), 1
-                if isinstance(vlit, slice):
-                    # Input was a SliceType variable, from which a slice was extracted.
-                    theslice = vlit
-                else:
-                    sslice = vlit
+                    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.")
-        if isinstance(sslice, int | np.integer | ScalarVariable):
-            return switch(lt(sslice, 0), (sslice + length), sslice), 1
+        cano = switch(lt(sslice, 0), (sslice + length), sslice)
+        return scalar_from_tensor(cano), 1
 
-    # At this point we have a slice object. Possibly symbolic.
+    # At this point we have a slice object. Possibly with symbolic inputs.
     assert isinstance(theslice, slice)
 
     def analyze(x):

tests/tensor/test_subtensor.py

@@ -16,7 +16,7 @@
 from pytensor.graph.op import get_test_value
 from pytensor.graph.rewriting.utils import is_same_graph
 from pytensor.printing import pprint
-from pytensor.scalar.basic import as_scalar
+from pytensor.scalar.basic import as_scalar, int16
 from pytensor.tensor import as_tensor, get_vector_length, vectorize
 from pytensor.tensor.blockwise import Blockwise
 from pytensor.tensor.elemwise import DimShuffle
@@ -68,7 +68,14 @@
     tensor5,
     vector,
 )
-from pytensor.tensor.type_other import NoneConst, SliceConstant, make_slice, slicetype
+from pytensor.tensor.type_other import (
+    NoneConst,
+    SliceConstant,
+    SliceType,
+    as_symbolic_slice,
+    make_slice,
+    slicetype,
+)
 from tests import unittest_tools as utt
 from tests.tensor.utils import inplace_func, integers_ranged, random
 
@@ -105,31 +112,51 @@ def test_as_index_literal():
 
 
 class TestGetCanonicalFormSlice:
-    def test_none_constant(self):
-        slc, rev = get_canonical_form_slice(NoneConst, 5)
-        assert isinstance(slc, slice)
-        assert slc.start == 0
-        assert slc.stop == 5
-        assert slc.step == 1
-        assert rev == 1
+    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)
 
     def test_scalar_constant(self):
         a = as_scalar(0)
         length = lscalar()
         res = get_canonical_form_slice(a, length)
-        assert res[0].owner.op == ptb.switch
+        assert isinstance(res[0].owner.op, ptb.ScalarFromTensor)
         assert res[1] == 1
 
     def test_tensor_constant(self):
         a = as_tensor(0)
         length = lscalar()
         res = get_canonical_form_slice(a, length)
-        assert res[0].owner.op == ptb.switch
+        assert isinstance(res[0].owner.op, ptb.ScalarFromTensor)
+        assert res[1] == 1
+
+    def test_symbolic_scalar(self):
+        a = int16()
+        length = lscalar()
+        res = get_canonical_form_slice(a, length)
+        assert res[0].owner.op, ptb.switch
+        assert res[1] == 1
+
+    def test_symbolic_tensor(self):
+        a = lscalar()
+        length = lscalar()
+        res = get_canonical_form_slice(a, length)
+        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 = make_slice(slice(3, 7, 2))
+        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