Fix rebase

Author: jaanerik

pytensor/link/numba/dispatch/subtensor.py

@@ -244,37 +244,7 @@ def numba_funcify_AdvancedSubtensor(op, node, **kwargs):
     else:
         tensor_inputs = node.inputs[2:]
 
-    adv_idxs = [
-        {
-            "axis": i,
-            "dtype": idx.type.dtype,
-            "bcast": idx.type.broadcastable,
-            "ndim": idx.type.ndim,
-        }
-        for i, idx in enumerate(idxs)
-        if isinstance(idx.type, TensorType)
-    ]
-
     # Reconstruct indexing information from idx_list and tensor inputs
-#   basic_idxs = []
-#   adv_idxs = []
-#   input_idx = 0
-#   
-#   for i, entry in enumerate(op.idx_list):
-#       if isinstance(entry, slice):
-#           # Basic slice index
-#           basic_idxs.append(entry)
-#       elif isinstance(entry, Type):
-#           # Advanced tensor index
-#           if input_idx < len(tensor_inputs):
-#               idx_input = tensor_inputs[input_idx]
-#               adv_idxs.append({
-#                   "axis": i,
-#                   "dtype": idx_input.type.dtype,
-#                   "bcast": idx_input.type.broadcastable,
-#                   "ndim": idx_input.type.ndim,
-#               })
-#               input_idx += 1
     basic_idxs = []
     adv_idxs = []
     input_idx = 0
@@ -313,7 +283,7 @@ def numba_funcify_AdvancedSubtensor(op, node, **kwargs):
         and len(adv_idxs) >= 1
         and all(adv_idx["dtype"] != "bool" for adv_idx in adv_idxs)
         # Implementation does not support newaxis
-        and not any(isinstance(idx.type, NoneTypeT) for idx in idxs)
+        and not any(isinstance(idx.type, NoneTypeT) for idx in tensor_inputs)
     ):
         return vector_integer_advanced_indexing(op, node, **kwargs)
 

pytensor/tensor/rewriting/subtensor.py

@@ -73,7 +73,6 @@
     IncSubtensor,
     Subtensor,
     advanced_inc_subtensor1,
-    advanced_subtensor,
     advanced_subtensor1,
     as_index_constant,
     get_canonical_form_slice,
@@ -83,7 +82,7 @@
     inc_subtensor,
     indices_from_subtensor,
 )
-from pytensor.tensor.type import TensorType
+from pytensor.tensor.type import TensorType, integer_dtypes
 from pytensor.tensor.type_other import NoneTypeT, SliceType
 from pytensor.tensor.variable import TensorConstant, TensorVariable
 
@@ -265,6 +264,7 @@ def local_AdvancedIncSubtensor_to_AdvancedIncSubtensor1(fgraph, node):
     """
 
     if type(node.op) is not AdvancedIncSubtensor:
+        # Don't apply to subclasses
         return
 
     if node.op.ignore_duplicates:
@@ -1321,7 +1321,9 @@ def local_useless_inc_subtensor_alloc(fgraph, node):
             if isinstance(node.op, IncSubtensor):
                 xi = Subtensor(node.op.idx_list)(x, *i)
             elif isinstance(node.op, AdvancedIncSubtensor):
-                xi = advanced_subtensor(x, *i)
+                # Use the same idx_list as the original operation to ensure correct shape
+                op = AdvancedSubtensor(node.op.idx_list)
+                xi = op.make_node(x, *i).outputs[0]
             elif isinstance(node.op, AdvancedIncSubtensor1):
                 xi = advanced_subtensor1(x, *i)
             else:
@@ -1771,10 +1773,11 @@ def local_blockwise_inc_subtensor(fgraph, node):
 
 
 @node_rewriter(tracks=[AdvancedSubtensor, AdvancedIncSubtensor])
-def bool_idx_to_nonzero(fgraph, node):
-    """Convert boolean indexing into equivalent vector boolean index, supported by our dispatch
+def ravel_multidimensional_bool_idx(fgraph, node):
+    """Convert multidimensional boolean indexing into equivalent vector boolean index, supported by Numba
 
-    x[1:, eye(3, dtype=bool), 1:] -> x[1:, *eye(3).nonzero()]
+    x[eye(3, dtype=bool)] -> x.ravel()[eye(3).ravel()]
+    x[eye(3, dtype=bool)].set(y) -> x.ravel()[eye(3).ravel()].set(y).reshape(x.shape)
     """
 
     if isinstance(node.op, AdvancedSubtensor):
@@ -1787,26 +1790,53 @@ def bool_idx_to_nonzero(fgraph, node):
     # Reconstruct indices from idx_list and tensor inputs
     idxs = indices_from_subtensor(tensor_inputs, node.op.idx_list)
 
-    bool_pos = {
-        i
+    if any(
+        (
+            (isinstance(idx.type, TensorType) and idx.type.dtype in integer_dtypes)
+            or isinstance(idx.type, NoneTypeT)
+        )
+        for idx in idxs
+    ):
+        # Get out if there are any other advanced indexes or np.newaxis
+        return None
+
+    bool_idxs = [
+        (i, idx)
         for i, idx in enumerate(idxs)
         if (isinstance(idx.type, TensorType) and idx.dtype == "bool")
-    }
+    ]
 
-    if not bool_pos:
+    if len(bool_idxs) != 1:
+        # Get out if there are no or multiple boolean idxs
+        return None
+    [(bool_idx_pos, bool_idx)] = bool_idxs
+    bool_idx_ndim = bool_idx.type.ndim
+    if bool_idx.type.ndim < 2:
+        # No need to do anything if it's a vector or scalar, as it's already supported by Numba
         return None
 
-    new_idxs = []
-    for i, idx in enumerate(idxs):
-        if i in bool_pos:
-            new_idxs.extend(idx.nonzero())
-        else:
-            new_idxs.append(idx)
+    x_shape = x.shape
+    raveled_x = x.reshape(
+        (*x_shape[:bool_idx_pos], -1, *x_shape[bool_idx_pos + bool_idx_ndim :])
+    )
+
+    raveled_bool_idx = bool_idx.ravel()
+    new_idxs = list(idxs)
+    new_idxs[bool_idx_pos] = raveled_bool_idx
 
     if isinstance(node.op, AdvancedSubtensor):
-        new_out = node.op(x, *new_idxs)
+        new_out = raveled_x[tuple(new_idxs)]
     else:
-        new_out = node.op(x, y, *new_idxs)
+        sub = raveled_x[tuple(new_idxs)]
+        new_out = inc_subtensor(
+            sub,
+            y,
+            set_instead_of_inc=node.op.set_instead_of_inc,
+            ignore_duplicates=node.op.ignore_duplicates,
+            inplace=node.op.inplace,
+        )
+        new_out = new_out.reshape(x_shape)
+
     return [copy_stack_trace(node.outputs[0], new_out)]
 
 
@@ -1941,10 +1971,16 @@ def ravel_multidimensional_int_idx(fgraph, node):
 
 
 optdb["specialize"].register(
-    bool_idx_to_nonzero.__name__,
-    bool_idx_to_nonzero,
+    ravel_multidimensional_bool_idx.__name__,
+    ravel_multidimensional_bool_idx,
+    "numba",
+    use_db_name_as_tag=False,  # Not included if only "specialize" is requested
+)
+
+optdb["specialize"].register(
+    ravel_multidimensional_int_idx.__name__,
+    ravel_multidimensional_int_idx,
     "numba",
-    "shape_unsafe",  # It can mask invalid mask sizes
     use_db_name_as_tag=False,  # Not included if only "specialize" is requested
 )
 

pytensor/tensor/subtensor.py

@@ -922,11 +922,12 @@ def __init__(self, idx_list=None):
 
     def _normalize_idx_list_for_hash(self):
         """Normalize idx_list for hash and equality comparison."""
-        if self.idx_list is None:
+        idx_list = getattr(self, "idx_list", None)
+        if idx_list is None:
             return None
 
         msg = []
-        for entry in self.idx_list:
+        for entry in idx_list:
             if isinstance(entry, slice):
                 msg.append((entry.start, entry.stop, entry.step))
             else:
@@ -2812,13 +2813,6 @@ def make_node(self, x, *inputs):
         advanced_indices = []
         adv_group_axis = None
         last_adv_group_axis = None
-        if new_axes: #not defined?
-            expanded_x_shape_list = list(x.type.shape)
-            for new_axis in new_axes:
-                expanded_x_shape_list.insert(new_axis, 1)
-            expanded_x_shape = tuple(expanded_x_shape_list)
-        else:
-            expanded_x_shape = x.type.shape
         for i, (idx, dim_length) in enumerate(
             zip_longest(explicit_indices, x.type.shape, fillvalue=slice(None))
         ):

tests/tensor/rewriting/test_basic.py

@@ -468,6 +468,63 @@ def test_incsubtensor(self):
         assert check_stack_trace(f1, ops_to_check="last")
         assert check_stack_trace(f2, ops_to_check="last")
 
+    def test_advanced_inc_subtensor_shape_inference_bug(self):
+        """
+        Test for bug in local_useless_inc_subtensor_alloc where advanced_subtensor
+        was called instead of using the original op's idx_list, causing incorrect
+        shape inference and AssertionError.
+
+        The bug occurred when advanced_subtensor(x, *i) tried to reconstruct
+        idx_list from inputs, leading to wrong shape for xi. This caused the
+        Assert condition checking shape compatibility to fail at runtime with:
+        AssertionError: `x[i]` and `y` do not have the same shape.
+
+        This test reproduces the bug by using a scenario where the shape
+        comparison would fail if xi has the wrong shape due to incorrect
+        idx_list reconstruction.
+        """
+        # Use vector with matrix indices - this creates AdvancedIncSubtensor
+        # The key is that when advanced_subtensor tries to reconstruct idx_list,
+        # it may get it wrong, causing xi to have incorrect shape
+        x = vector("x")
+        y = scalar("y")
+        i = matrix(
+            "i", dtype="int64"
+        )  # 2D indices for 1D array -> AdvancedIncSubtensor
+
+        # Create AdvancedIncSubtensor with Alloc
+        # When i is (n, m), i.shape is (n, m), so alloc creates shape (n, m)
+        # But x[i] where i is (n, m) creates shape (n, m) as well
+        # The bug would cause xi to have wrong shape, making the Assert fail
+        z = advanced_inc_subtensor(x, pt.alloc(y, *i.shape), i)
+
+        # Compile - this should not raise AssertionError during execution
+        # With the buggy code (using advanced_subtensor), this raises:
+        # AssertionError: `x[i]` and `y` do not have the same shape.
+        f = function([x, i, y], z, mode=self.mode)
+
+        # Test with actual values
+        x_value = np.random.standard_normal(10).astype(config.floatX)
+        y_value = np.random.standard_normal()
+        i_value = self.rng.integers(0, 10, size=(3, 2))
+
+        # This should execute without AssertionError
+        # With the buggy code (using advanced_subtensor), this would raise:
+        # AssertionError: `x[i]` and `y` do not have the same shape.
+        result = f(x_value, i_value, y_value)
+
+        # Verify basic properties
+        # The main point of this test is that it doesn't raise AssertionError
+        # advanced_inc_subtensor modifies x in place and returns it
+        assert result.shape == x_value.shape, "Result should have same shape as input"
+        assert not np.array_equal(result, x_value), "Result should be modified"
+
+        # Verify the rewrite was applied (Alloc should be removed)
+        topo = f.maker.fgraph.toposort()
+        assert len([n for n in topo if isinstance(n.op, Alloc)]) == 0, (
+            "Alloc should have been removed by the rewrite"
+        )
+
 
 class TestUselessCheckAndRaise:
     def test_basic(self):

tests/tensor/rewriting/test_elemwise.py

@@ -1642,9 +1642,9 @@ def test_InplaceElemwiseOptimizer_bug():
     # with config.change_flags(tensor__insert_inplace_optimizer_validate_nb=10):
     rewrite_graph(fgraph, include=("inplace",))
 
-    pytensor.config.tensor__insert_inplace_optimizer_validate_nb = 1
-    with pytest.warns(
-        FutureWarning,
-        match="tensor__insert_inplace_optimizer_validate_nb config is deprecated",
-    ):
-        rewrite_graph(fgraph, include=("inplace",))
+    with config.change_flags(tensor__insert_inplace_optimizer_validate_nb=1):
+        with pytest.warns(
+            FutureWarning,
+            match="tensor__insert_inplace_optimizer_validate_nb config is deprecated",
+        ):
+            rewrite_graph(fgraph, include=("inplace",))

tests/tensor/test_basic.py

@@ -705,7 +705,7 @@ def test_masked_array_not_implemented(
 
 
 def check_alloc_runtime_broadcast(mode):
-    """Check we emmit a clear error when runtime broadcasting would occur according to Numpy rules."""
+    """Check we emit a clear error when runtime broadcasting would occur according to Numpy rules."""
     floatX = config.floatX
     x_v = vector("x", shape=(None,))
 

tests/tensor/test_subtensor.py

@@ -11,7 +11,7 @@
 import pytensor
 import pytensor.scalar as scal
 import pytensor.tensor.basic as ptb
-from pytensor import config, function, shared
+from pytensor import function, shared
 from pytensor.compile import DeepCopyOp
 from pytensor.compile.io import In
 from pytensor.compile.mode import Mode, get_default_mode
@@ -3301,6 +3301,33 @@ def test_slice_at_axis():
     assert x_sliced.type.shape == (3, 1, 5)
 
 
+def test_advanced_inc_subtensor1_failure():
+    # Shapes from the failure log
+    N = 500
+    TotalCols = 7
+    OrderedCols = 5
+    UnorderedCols = 2
+
+    oinds_val = [1, 2, 3, 5, 6]
+    uoinds_val = [0, 4]
+
+    y_ordered = matrix("y_ordered")
+    y_unordered = matrix("y_unordered")
+
+    fodds_init = ptb.empty((N, TotalCols))
+
+    fodds_step1 = set_subtensor(fodds_init[:, uoinds_val], y_unordered)
+    fodds_step2 = set_subtensor(fodds_step1[:, oinds_val], y_ordered)
+
+    f = pytensor.function([y_unordered, y_ordered], fodds_step2)
+    # assert any("AdvancedIncSubtensor1" in str(node) for node in f.maker.fgraph.toposort())
+
+    y_u_data = np.random.randn(N, UnorderedCols).astype(np.float64)
+    y_o_data = np.random.randn(N, OrderedCols).astype(np.float64)
+    res = f(y_u_data, y_o_data)
+    assert res.shape == (N, TotalCols)
+
+
 @pytest.mark.parametrize(
     "size", [(3,), (3, 3), (3, 5, 5)], ids=["1d", "2d square", "3d square"]
 )