Implement simpler idx_list

Author: jaanerik

pytensor/graph/destroyhandler.py

@@ -771,9 +771,9 @@ def orderings(self, fgraph, ordered=True):
                     }
                     tolerated.add(destroyed_idx)
                     tolerate_aliased = getattr(
-                        app.op, "destroyhandler_tolerate_aliased", []
+                        app.op, "destroyhandler_tolerate_aliased", ()
                     )
-                    assert isinstance(tolerate_aliased, list)
+                    assert isinstance(tolerate_aliased, list | tuple)
                     ignored = {
                         idx1 for idx0, idx1 in tolerate_aliased if idx0 == destroyed_idx
                     }

pytensor/link/numba/dispatch/subtensor.py

@@ -10,7 +10,7 @@
 from numba.core.pythonapi import box
 
 import pytensor.link.numba.dispatch.basic as numba_basic
-from pytensor.graph import Type, Variable
+from pytensor.graph import Variable
 from pytensor.link.numba.cache import (
     compile_numba_function_src,
 )
@@ -29,6 +29,7 @@
     AdvancedSubtensor1,
     IncSubtensor,
     Subtensor,
+    _is_position,
     indices_from_subtensor,
 )
 from pytensor.tensor.type_other import MakeSlice, NoneTypeT
@@ -158,7 +159,7 @@ def numba_funcify_default_subtensor(op, node, **kwargs):
     """Create a Python function that assembles and uses an index on an array."""
 
     def convert_indices(indices_iterator, entry):
-        if hasattr(indices_iterator, "__next__") and isinstance(entry, Type):
+        if hasattr(indices_iterator, "__next__") and _is_position(entry):
             name, var = next(indices_iterator)
             if var.ndim == 0 and isinstance(var.type, TensorType):
                 return f"{name}.item()"
@@ -171,8 +172,6 @@ def convert_indices(indices_iterator, entry):
             )
         elif isinstance(entry, type(None)):
             return "None"
-        elif isinstance(entry, (int, np.integer)):
-            return str(entry)
         else:
             raise ValueError(f"Unknown index type: {entry}")
 

pytensor/tensor/basic.py

@@ -29,7 +29,7 @@
 from pytensor.graph.op import Op
 from pytensor.graph.replace import _vectorize_node
 from pytensor.graph.rewriting.db import EquilibriumDB
-from pytensor.graph.type import HasShape, Type
+from pytensor.graph.type import HasShape
 from pytensor.link.c.op import COp
 from pytensor.link.c.params_type import ParamsType
 from pytensor.printing import Printer, min_informative_str, pprint, set_precedence
@@ -300,7 +300,7 @@ def _get_underlying_scalar_constant_value(
     """
     from pytensor.compile.ops import DeepCopyOp, OutputGuard
     from pytensor.sparse import CSM
-    from pytensor.tensor.subtensor import Subtensor
+    from pytensor.tensor.subtensor import Subtensor, _is_position
 
     v = orig_v
     while True:
@@ -433,7 +433,7 @@ def _get_underlying_scalar_constant_value(
                         var.ndim == 1 for var in v.owner.inputs[0].owner.inputs[1:]
                     ):
                         idx = v.owner.op.idx_list[0]
-                        if isinstance(idx, Type):
+                        if _is_position(idx):
                             idx = _get_underlying_scalar_constant_value(
                                 v.owner.inputs[1], max_recur=max_recur
                             )
@@ -467,7 +467,7 @@ def _get_underlying_scalar_constant_value(
                     and len(v.owner.op.idx_list) == 1
                 ):
                     idx = v.owner.op.idx_list[0]
-                    if isinstance(idx, Type):
+                    if _is_position(idx):
                         idx = _get_underlying_scalar_constant_value(
                             v.owner.inputs[1], max_recur=max_recur
                         )
@@ -488,7 +488,7 @@ def _get_underlying_scalar_constant_value(
                     op = owner.op
                     idx_list = op.idx_list
                     idx = idx_list[0]
-                    if isinstance(idx, Type):
+                    if _is_position(idx):
                         idx = _get_underlying_scalar_constant_value(
                             owner.inputs[1], max_recur=max_recur
                         )

pytensor/tensor/rewriting/shape.py

@@ -17,7 +17,6 @@
 )
 from pytensor.graph.traversal import ancestors
 from pytensor.graph.utils import InconsistencyError, get_variable_trace_string
-from pytensor.scalar import ScalarType
 from pytensor.tensor.basic import (
     MakeVector,
     as_tensor_variable,
@@ -45,7 +44,7 @@
     SpecifyShape,
     specify_shape,
 )
-from pytensor.tensor.subtensor import Subtensor, get_idx_list
+from pytensor.tensor.subtensor import Subtensor, _is_position, get_idx_list
 from pytensor.tensor.type import TensorType, discrete_dtypes, integer_dtypes
 from pytensor.tensor.type_other import NoneTypeT
 from pytensor.tensor.variable import TensorVariable
@@ -845,13 +844,16 @@ def _is_shape_i_of_x(
     if isinstance(var.owner.op, Shape_i):
         return (var.owner.op.i == i) and (var.owner.inputs[0] == x)  # type: ignore
 
-    # Match Subtensor((ScalarType,))(Shape(input), i)
+    # Match Subtensor((int,))(Shape(input), i) - single integer index into shape
     if isinstance(var.owner.op, Subtensor):
+        idx_entry = (
+            var.owner.op.idx_list[0] if len(var.owner.op.idx_list) == 1 else None
+        )
         return (
             # Check we have integer indexing operation
             # (and not slice or multiple indexing)
             len(var.owner.op.idx_list) == 1
-            and isinstance(var.owner.op.idx_list[0], ScalarType)
+            and _is_position(idx_entry)
             # Check we are indexing on the shape of x
             and var.owner.inputs[0].owner is not None
             and isinstance(var.owner.inputs[0].owner.op, Shape)

pytensor/tensor/rewriting/subtensor.py

@@ -6,7 +6,7 @@
 import pytensor
 from pytensor import compile
 from pytensor.compile import optdb
-from pytensor.graph.basic import Constant, Variable
+from pytensor.graph.basic import Constant, Variable, equal_computations
 from pytensor.graph.rewriting.basic import (
     WalkingGraphRewriter,
     copy_stack_trace,
@@ -15,7 +15,7 @@
     node_rewriter,
 )
 from pytensor.raise_op import Assert
-from pytensor.scalar import Add, ScalarConstant, ScalarType
+from pytensor.scalar import Add, ScalarConstant
 from pytensor.scalar import constant as scalar_constant
 from pytensor.tensor.basic import (
     Alloc,
@@ -72,6 +72,7 @@
     AdvancedSubtensor1,
     IncSubtensor,
     Subtensor,
+    _is_position,
     advanced_inc_subtensor1,
     advanced_subtensor1,
     as_index_constant,
@@ -480,9 +481,8 @@ def local_subtensor_remove_broadcastable_index(fgraph, node):
     remove_dim = []
     node_inputs_idx = 1
     for dim, elem in enumerate(idx):
-        if isinstance(elem, ScalarType):
-            # The idx is a ScalarType, ie a Type. This means the actual index
-            # is contained in node.inputs[1]
+        if _is_position(elem):
+            # The idx is a integer position.
             dim_index = node.inputs[node_inputs_idx]
             if isinstance(dim_index, ScalarConstant):
                 dim_index = dim_index.value
@@ -494,9 +494,6 @@ def local_subtensor_remove_broadcastable_index(fgraph, node):
         elif isinstance(elem, slice):
             if elem != slice(None):
                 return
-        elif isinstance(elem, int | np.integer):
-            if elem in (0, -1) and node.inputs[0].broadcastable[dim]:
-                remove_dim.append(dim)
         else:
             raise TypeError("case not expected")
 
@@ -508,6 +505,39 @@ def local_subtensor_remove_broadcastable_index(fgraph, node):
         return [node.inputs[0].dimshuffle(tuple(remain_dim))]
 
 
+def _idx_list_struct_equal(idx_list1, idx_list2):
+    """Check if two idx_lists have the same structure.
+
+    Positions (integers) are treated as equivalent regardless of value,
+    since positions are relative to each Op's inputs.
+    """
+    if len(idx_list1) != len(idx_list2):
+        return False
+
+    def normalize_entry(entry):
+        if isinstance(entry, int) and not isinstance(entry, bool):
+            return "POS"  # All positions are equivalent
+        elif isinstance(entry, slice):
+            return (
+                "POS"
+                if isinstance(entry.start, int) and not isinstance(entry.start, bool)
+                else entry.start,
+                "POS"
+                if isinstance(entry.stop, int) and not isinstance(entry.stop, bool)
+                else entry.stop,
+                "POS"
+                if isinstance(entry.step, int) and not isinstance(entry.step, bool)
+                else entry.step,
+            )
+        else:
+            return entry
+
+    for e1, e2 in zip(idx_list1, idx_list2):
+        if normalize_entry(e1) != normalize_entry(e2):
+            return False
+    return True
+
+
 @register_specialize
 @register_canonicalize
 @node_rewriter([Subtensor])
@@ -523,9 +553,17 @@ def local_subtensor_inc_subtensor(fgraph, node):
         if not x.owner.op.set_instead_of_inc:
             return
 
-        if x.owner.inputs[2:] == node.inputs[1:] and tuple(
-            x.owner.op.idx_list
-        ) == tuple(node.op.idx_list):
+        # Check structural equality of idx_lists and semantic equality of inputs
+        inc_inputs = x.owner.inputs[2:]
+        sub_inputs = node.inputs[1:]
+
+        if (
+            len(inc_inputs) == len(sub_inputs)
+            and _idx_list_struct_equal(x.owner.op.idx_list, node.op.idx_list)
+            and all(
+                equal_computations([a], [b]) for a, b in zip(inc_inputs, sub_inputs)
+            )
+        ):
             out = node.outputs[0]
             y = x.owner.inputs[1]
             # If the dtypes differ, cast y into x.dtype

pytensor/tensor/rewriting/subtensor_lift.py

@@ -8,7 +8,6 @@
 from pytensor.compile import optdb
 from pytensor.graph import Constant, FunctionGraph, node_rewriter, vectorize_graph
 from pytensor.graph.rewriting.basic import NodeRewriter, copy_stack_trace
-from pytensor.scalar import basic as ps
 from pytensor.tensor.basic import (
     Alloc,
     Join,
@@ -42,6 +41,7 @@
     AdvancedSubtensor,
     AdvancedSubtensor1,
     Subtensor,
+    _is_position,
     _non_consecutive_adv_indexing,
     as_index_literal,
     get_canonical_form_slice,
@@ -702,13 +702,13 @@ def local_subtensor_make_vector(fgraph, node):
 
         (idx,) = idxs
 
-        if isinstance(idx, ps.ScalarType | TensorType):
-            old_idx, idx = idx, node.inputs[1]
-            assert idx.type.is_super(old_idx)
+        if _is_position(idx):
+            # idx is an integer position - get the actual index value from inputs
+            idx = node.inputs[1]
     elif isinstance(node.op, AdvancedSubtensor1):
         idx = node.inputs[1]
 
-    if isinstance(idx, int | np.integer):
+    if False:  # isinstance(idx, int | np.integer) - disabled, positions handled above
         return [x.owner.inputs[idx]]
     elif isinstance(idx, Variable):
         if idx.ndim == 0: