Don't run local_uint_constant_indices in C/python backend

Indices are always cast to int64 by the underlying methods.

Also don't run in specialize, to reduce number of passes. Other rewrites may introduce temporar indexing operations (such as x.shape[i]) which always default to int64, and it's useless to optimize immediately.

Author: ricardoV94

pytensor/compile/mode.py

@@ -489,7 +489,6 @@ def clone(self, link_kwargs=None, optimizer="", **kwargs):
             "BlasOpt",
             "fusion",
             "inplace",
-            "local_uint_constant_indices",
             "scan_save_mem_prealloc",
         ],
     ),

pytensor/tensor/rewriting/subtensor.py

@@ -5,7 +5,6 @@
 import numpy as np
 
 import pytensor
-import pytensor.scalar.basic as ps
 from pytensor import compile
 from pytensor.compile import optdb
 from pytensor.graph.basic import Constant, Variable
@@ -14,8 +13,11 @@
     copy_stack_trace,
     in2out,
     node_rewriter,
+    out2in,
 )
 from pytensor.raise_op import Assert
+from pytensor.scalar import Add, ScalarConstant, ScalarType
+from pytensor.scalar import constant as scalar_constant
 from pytensor.tensor.basic import (
     Alloc,
     Join,
@@ -31,6 +33,7 @@
     register_infer_shape,
     switch,
 )
+from pytensor.tensor.basic import constant as tensor_constant
 from pytensor.tensor.blockwise import Blockwise
 from pytensor.tensor.elemwise import Elemwise
 from pytensor.tensor.exceptions import NotScalarConstantError
@@ -588,11 +591,11 @@ def local_subtensor_remove_broadcastable_index(fgraph, node):
     remove_dim = []
     node_inputs_idx = 1
     for dim, elem in enumerate(idx):
-        if isinstance(elem, (ps.ScalarType)):
+        if isinstance(elem, ScalarType):
             # The idx is a ScalarType, ie a Type. This means the actual index
             # is contained in node.inputs[1]
             dim_index = node.inputs[node_inputs_idx]
-            if isinstance(dim_index, ps.ScalarConstant):
+            if isinstance(dim_index, ScalarConstant):
                 dim_index = dim_index.value
             if dim_index in (0, -1) and node.inputs[0].broadcastable[dim]:
                 remove_dim.append(dim)
@@ -770,7 +773,7 @@ def local_subtensor_make_vector(fgraph, node):
 
         (idx,) = idxs
 
-        if isinstance(idx, ps.ScalarType | TensorType):
+        if isinstance(idx, ScalarType | TensorType):
             old_idx, idx = idx, node.inputs[1]
             assert idx.type.is_super(old_idx)
     elif isinstance(node.op, AdvancedSubtensor1):
@@ -895,7 +898,7 @@ def local_set_to_inc_subtensor(fgraph, node):
         and node.op.set_instead_of_inc
         and node.inputs[1].owner
         and isinstance(node.inputs[1].owner.op, Elemwise)
-        and isinstance(node.inputs[1].owner.op.scalar_op, ps.Add)
+        and isinstance(node.inputs[1].owner.op.scalar_op, Add)
     ):
         addn = node.inputs[1].owner
         subn = None
@@ -1789,7 +1792,6 @@ def local_join_subtensors(fgraph, node):
                 return [merged_subtensors]
 
 
-@register_specialize
 @node_rewriter(
     [
         Subtensor,
@@ -1850,12 +1852,10 @@ def local_uint_constant_indices(fgraph, node):
         if dtype == index_val.dtype:
             continue
 
-        if index_val.ndim > 0:
-            new_index = pytensor.tensor.as_tensor_variable(
-                index_val.astype(dtype), dtype=dtype
-            )
+        if isinstance(index.type, TensorType):
+            new_index = tensor_constant(index_val.astype(dtype), dtype=dtype)
         else:
-            new_index = ps.constant(index_val.astype(dtype), dtype=dtype)
+            new_index = scalar_constant(index_val.astype(dtype), dtype=dtype)
 
         new_indices[i] = new_index
         has_new_index = True
@@ -1877,6 +1877,20 @@ def local_uint_constant_indices(fgraph, node):
     return [new_out]
 
 
+compile.optdb.register(
+    local_uint_constant_indices.__name__,
+    out2in(local_uint_constant_indices),
+    # Python / C backends always cast indices to int64 internally.
+    "numba",
+    "jax",
+    # After specialization and uncanonicalization
+    # Other rewrites don't worry about the dtype of the indices
+    # And can cause unnecessary passes of this optimization
+    # Such as x.shape[np.int(0)] -> x.shape[np.uint(0)]
+    position=4,
+)
+
+
 @register_canonicalize("shape_unsafe")
 @register_stabilize("shape_unsafe")
 @register_specialize("shape_unsafe")