Adapt to AdvancedIndexing changes

Author: ricardoV94

pymc/logprob/mixture.py

@@ -36,12 +36,11 @@
 
 from typing import cast
 
-import pytensor
 import pytensor.tensor as pt
 
 from pytensor.graph.basic import Apply, Constant, Variable
 from pytensor.graph.fg import FunctionGraph
-from pytensor.graph.op import Op, compute_test_value
+from pytensor.graph.op import Op
 from pytensor.graph.rewriting.basic import EquilibriumGraphRewriter, node_rewriter
 from pytensor.graph.traversal import ancestors
 from pytensor.ifelse import IfElse, ifelse
@@ -60,10 +59,10 @@
     AdvancedSubtensor1,
     as_index_literal,
     get_canonical_form_slice,
-    is_basic_idx,
+    unflatten_index_variables,
 )
 from pytensor.tensor.type import TensorType
-from pytensor.tensor.type_other import NoneConst, NoneTypeT, SliceType
+from pytensor.tensor.type_other import NoneConst, NoneTypeT
 from pytensor.tensor.variable import TensorVariable
 
 from pymc.logprob.abstract import (
@@ -89,12 +88,9 @@
 from pymc.pytensorf import constant_fold
 
 
-def is_newaxis(x):
-    return isinstance(x, type(None)) or isinstance(getattr(x, "type", None), NoneTypeT)
-
-
 def expand_indices(
-    indices: tuple[Variable | slice | None, ...], shape: tuple[TensorVariable]
+    indices: tuple[Variable | slice, ...],
+    shape: tuple[TensorVariable],
 ) -> tuple[TensorVariable]:
     """Convert basic and/or advanced indices into a single, broadcasted advanced indexing operation.
 
@@ -106,32 +102,34 @@ def expand_indices(
         The shape of the array being indexed.
 
     """
-    n_non_newaxis = sum(1 for idx in indices if not is_newaxis(idx))
-    n_missing_dims = len(shape) - n_non_newaxis
-    full_indices = list(indices) + [slice(None)] * n_missing_dims
+    # Make implicit slice(None) explicit
+    full_indices = list(indices) + [slice(None)] * (len(shape) - len(indices))
 
-    # We need to know if a "subspace" was generated by advanced indices
-    # bookending basic indices.  If so, we move the advanced indexing subspace
-    # to the "front" of the shape (i.e. left-most indices/last-most
-    # dimensions).
-    index_types = [is_basic_idx(idx) for idx in full_indices]
+    # We need to know if a "subspace" was generated by advanced indices bookending basic indices.
+    # If so, we move the advanced indexing subspace to the "front" of the shape
+    # (i.e. left-most indices/last-most dimensions).
+    # NOTE: Scalar integers are classified as "advanced" here even though they may be basic
+    # in a pure-basic-indexing context. This is safe because when all non-slice indices are
+    # scalars, n_subspace_dims=0 and the moved_subspace logic becomes a no-op.
+    is_basic_idx = [isinstance(idx, slice) for idx in full_indices]
 
     first_adv_idx = len(shape)
     try:
-        first_adv_idx = index_types.index(False)
-        first_bsc_after_adv_idx = index_types.index(True, first_adv_idx)
-        index_types.index(False, first_bsc_after_adv_idx)
+        first_adv_idx = is_basic_idx.index(False)
+        first_bsc_after_adv_idx = is_basic_idx.index(True, first_adv_idx)
+        # If this fails, we get to the except branch
+        is_basic_idx.index(False, first_bsc_after_adv_idx)
         moved_subspace = True
     except ValueError:
         moved_subspace = False
 
-    n_basic_indices = sum(index_types)
+    n_basic_indices = sum(is_basic_idx)
 
     # The number of dimensions in the subspace created by the advanced indices
     n_subspace_dims = max(
         (
             getattr(idx, "ndim", 0)
-            for idx, is_basic in zip(full_indices, index_types)
+            for idx, is_basic in zip(full_indices, is_basic_idx)
             if not is_basic
         ),
         default=0,
@@ -144,63 +142,53 @@ def expand_indices(
     shape_copy = list(shape)
     n_preceding_basics = 0
     for d, idx in enumerate(full_indices):
-        if not is_basic_idx(idx):
-            s = shape_copy.pop(0)
-
+        s = shape_copy.pop(0)
+        if not isinstance(idx, slice):
+            # Advanced indexing
             idx = pt.as_tensor(idx)
 
             if moved_subspace:
-                # The subspace generated by advanced indices appear as the
-                # upper dimensions in the "expanded" index space, so we need to
-                # add broadcast dimensions for the non-basic indices to the end
-                # of these advanced indices
-                expanded_idx = idx[(Ellipsis,) + (None,) * n_basic_indices]
+                # The subspace generated by advanced indices appear as the leftmost dimensions
+                # in the "expanded" index space, so we need to add broadcast dimensions
+                # for the non-basic indices to the end of these advanced indices
+                expanded_idx = pt.shape_padright(idx, n_basic_indices)
             else:
                 # In this case, we need to add broadcast dimensions for the
                 # basic indices that proceed and follow the group of advanced
                 # indices; otherwise, a contiguous group of advanced indices
                 # forms a broadcasted set of indices that are iterated over
                 # within the same subspace, which means that all their
                 # corresponding "expanded" indices have exactly the same shape.
-                expanded_idx = idx[(None,) * n_preceding_basics][
-                    (Ellipsis,) + (None,) * (n_basic_indices - n_preceding_basics)
-                ]
+                expanded_idx = pt.shape_padleft(idx, n_preceding_basics)
+                expanded_idx = pt.shape_padright(expanded_idx, n_basic_indices - n_preceding_basics)
         else:
-            if is_newaxis(idx):
-                n_preceding_basics += 1
-                continue
-
-            s = shape_copy.pop(0)
-
-            if isinstance(idx, slice) or isinstance(getattr(idx, "type", None), SliceType):
-                idx = as_index_literal(idx)
-                idx_slice, _ = get_canonical_form_slice(idx, s)
-                idx = pt.arange(idx_slice.start, idx_slice.stop, idx_slice.step)
+            # Basic indexing (slice)
+            idx = as_index_literal(idx)
+            idx_slice, _ = get_canonical_form_slice(idx, s)
+            idx = pt.arange(idx_slice.start, idx_slice.stop, idx_slice.step)
 
             if moved_subspace:
-                # Basic indices appear in the lower dimensions
-                # (i.e. right-most) in the output, and are preceded by
-                # the subspace generated by the advanced indices.
-                expanded_idx = idx[(None,) * (n_subspace_dims + n_preceding_basics)][
-                    (Ellipsis,) + (None,) * (n_basic_indices - n_preceding_basics - 1)
-                ]
+                # Basic indices appear in the rightmost dimensions in the output,
+                # and are preceded by the subspace generated by the advanced indices.
+                expanded_idx = pt.shape_padleft(idx, n_subspace_dims + n_preceding_basics)
+                expanded_idx = pt.shape_padright(
+                    expanded_idx, n_basic_indices - n_preceding_basics - 1
+                )
             else:
                 # In this case, we need to know when the basic indices have
-                # moved past the contiguous group of advanced indices (in the
-                # "expanded" index space), so that we can properly pad those
-                # dimensions in this basic index's shape.
+                # moved past the contiguous group of advanced indices
+                # (in the "expanded" index space), so that we can properly
+                # pad those  dimensions in this basic index's shape.
                 # Don't forget that a single advanced index can introduce an
                 # arbitrary number of dimensions to the expanded index space.
 
-                # If we're currently at a basic index that's past the first
-                # advanced index, then we're necessarily past the group of
-                # advanced indices.
+                # If we're currently at a basic index that's past the first advanced index,
+                # then we're necessarily past the group of  advanced indices.
                 n_preceding_dims = (
                     n_subspace_dims if d > first_adv_idx else 0
                 ) + n_preceding_basics
-                expanded_idx = idx[(None,) * n_preceding_dims][
-                    (Ellipsis,) + (None,) * (n_output_dims - n_preceding_dims - 1)
-                ]
+                expanded_idx = pt.shape_padleft(idx, n_preceding_dims)
+                expanded_idx = pt.shape_padright(expanded_idx, n_output_dims - n_preceding_dims - 1)
 
             n_preceding_basics += 1
 
@@ -227,13 +215,14 @@ def rv_pull_down(x: TensorVariable) -> TensorVariable:
 class MixtureRV(MeasurableOp, Op):
     """A placeholder used to specify a log-likelihood for a mixture sub-graph."""
 
-    __props__ = ("indices_end_idx", "out_dtype", "out_broadcastable")
+    __props__ = ("indices_end_idx", "out_dtype", "out_broadcastable", "idx_list")
 
-    def __init__(self, indices_end_idx, out_dtype, out_broadcastable):
+    def __init__(self, indices_end_idx, out_dtype, out_broadcastable, idx_list):
         super().__init__()
         self.indices_end_idx = indices_end_idx
         self.out_dtype = out_dtype
         self.out_broadcastable = out_broadcastable
+        self.idx_list = tuple(idx_list)
 
     def make_node(self, *inputs):
         return Apply(self, list(inputs), [TensorType(self.out_dtype, self.out_broadcastable)()])
@@ -285,6 +274,7 @@ def find_measurable_index_mixture(fgraph, node):
     created for each ``i`` in ``enumerate(mixture_comps)``.
     """
     mixing_indices = node.inputs[1:]
+    idx_list = node.op.idx_list
 
     # TODO: Add check / test case for Advanced Boolean indexing
     if isinstance(node.op, AdvancedSubtensor | AdvancedSubtensor1):
@@ -313,18 +303,9 @@ def find_measurable_index_mixture(fgraph, node):
         1 + len(mixing_indices),
         old_mixture_rv.dtype,
         old_mixture_rv.broadcastable,
+        idx_list=idx_list,
     )
-    new_node = mix_op.make_node(*([join_axis, *mixing_indices, *mixture_rvs]))
-
-    new_mixture_rv = new_node.default_output()
-
-    if pytensor.config.compute_test_value != "off":
-        # We can't use `MixtureRV` to compute a test value; instead, we'll use
-        # the original node's test value.
-        if not hasattr(old_mixture_rv.tag, "test_value"):
-            compute_test_value(node)
-
-        new_mixture_rv.tag.test_value = old_mixture_rv.tag.test_value
+    new_mixture_rv = mix_op(join_axis, *mixing_indices, *mixture_rvs)
 
     return [new_mixture_rv]
 
@@ -334,16 +315,17 @@ def logprob_MixtureRV(op, values, *inputs: TensorVariable | slice | None, name=N
     (value,) = values
 
     join_axis = cast(Variable, inputs[0])
-    indices = cast(TensorVariable, inputs[1 : op.indices_end_idx])
+    flat_indices = inputs[1 : op.indices_end_idx]
     comp_rvs = cast(TensorVariable, inputs[op.indices_end_idx :])
 
-    assert len(indices) > 0
+    assert len(flat_indices) > 0
 
-    if len(indices) > 1 or indices[0].ndim > 0:
+    # Use flat_indices (tensor variables only) for the scalar vs advanced check.
+    # Slices in idx_list don't indicate advanced indexing.
+    if len(flat_indices) > 1 or flat_indices[0].ndim > 0:
         if isinstance(join_axis.type, NoneTypeT):
             # `join_axis` will be `NoneConst` if the "join" was a `MakeVector`
-            # (i.e. scalar measurable variables were combined to make a
-            # vector).
+            # (i.e. scalar measurable variables were combined to make a vector).
             # Since some form of advanced indexing is necessarily occurring, we
             # need to reformat the MakeVector arguments so that they fit the
             # `Join` format expected by the logic below.
@@ -354,7 +336,9 @@ def logprob_MixtureRV(op, values, *inputs: TensorVariable | slice | None, name=N
             join_axis_val = constant_fold((join_axis,))[0].item()
             original_shape = shape_tuple(comp_rvs[0])
 
-        bcast_indices = expand_indices(indices, original_shape)
+        # Reconstruct full indices
+        full_indices = unflatten_index_variables(list(flat_indices), op.idx_list)
+        bcast_indices = expand_indices(full_indices, original_shape)
 
         logp_val = pt.empty(bcast_indices[0].shape)
 
@@ -393,7 +377,7 @@ def logprob_MixtureRV(op, values, *inputs: TensorVariable | slice | None, name=N
             if join_axis_val is not None:
                 comp_logp = pt.squeeze(comp_logp, axis=join_axis_val)
             logp_val += ifelse(
-                pt.eq(indices[0], i),
+                pt.eq(flat_indices[0], i),
                 comp_logp,
                 pt.zeros_like(comp_logp),
             )

tests/logprob/test_mixture.py

@@ -827,61 +827,6 @@ def test_expand_indices_single_indices(A_parts, indices):
     assert np.array_equal(res, exp_res)
 
 
-@pytest.mark.parametrize(
-    "A_parts, indices",
-    [
-        (
-            (
-                np.random.normal(size=(4, 3)),
-                np.random.normal(size=(4, 3)),
-                np.random.normal(size=(4, 3)),
-            ),
-            (None,),
-        ),
-        (
-            (
-                np.random.normal(size=(4, 3)),
-                np.random.normal(size=(4, 3)),
-                np.random.normal(size=(4, 3)),
-            ),
-            (None, None, None),
-        ),
-        (
-            (
-                np.random.normal(size=(4, 3)),
-                np.random.normal(size=(4, 3)),
-                np.random.normal(size=(4, 3)),
-            ),
-            (None, 1, None, 0, None),
-        ),
-        (
-            (
-                np.random.normal(size=(4, 3)),
-                np.random.normal(size=(4, 3)),
-                np.random.normal(size=(4, 3)),
-            ),
-            (slice(2, 3), None, 1, None, 0, None),
-        ),
-        (
-            (
-                np.random.normal(size=(4, 3)),
-                np.random.normal(size=(4, 3)),
-                np.random.normal(size=(4, 3)),
-            ),
-            (slice(2, 3), None, 1, 0, None),
-        ),
-    ],
-)
-def test_expand_indices_newaxis(A_parts, indices):
-    A = pt.stack(A_parts)
-    at_indices = [as_index_constant(idx) for idx in indices]
-    full_indices = expand_indices(at_indices, shape_tuple(A))
-    assert len(full_indices) == A.ndim
-    exp_res = A[indices].eval()
-    res = A[full_indices].eval()
-    assert np.array_equal(res, exp_res)
-
-
 def test_mixture_with_DiracDelta():
     srng = pt.random.RandomStream(29833)