Workaround buggy AdvancedIndexing Mixture logprob

Now that Dimshuffle lift broadcasts both the parameters and the size, the AdvancedIndexing logprob fails most of the times, even though these are valid graphs.

All but one of the failing cases can be helped by introducing the local_rv_size_lift rewrite.

Author: ricardoV94

pymc/logprob/mixture.py

@@ -53,6 +53,7 @@
 from pytensor.tensor.elemwise import Elemwise
 from pytensor.tensor.random.rewriting import (
     local_dimshuffle_rv_lift,
+    local_rv_size_lift,
     local_subtensor_rv_lift,
 )
 from pytensor.tensor.shape import shape_tuple
@@ -210,6 +211,7 @@ def rv_pull_down(x: TensorVariable, dont_touch_vars=None) -> TensorVariable:
     return pre_greedy_node_rewriter(
         fgraph,
         [
+            local_rv_size_lift,
             local_dimshuffle_rv_lift,
             local_subtensor_rv_lift,
             naive_bcast_rv_lift,

pymc/tests/logprob/test_mixture.py

@@ -324,7 +324,7 @@ def test_hetero_mixture_binomial(p_val, size):
             0,
         ),
         # Same as before but with degenerate vector parameters
-        pytest.param(
+        (
             (
                 np.array([0], dtype=pytensor.config.floatX),
                 np.array(1, dtype=pytensor.config.floatX),
@@ -342,7 +342,6 @@ def test_hetero_mixture_binomial(p_val, size):
             (2,),
             (),
             0,
-            marks=pytest.mark.xfail(IndexError, reason="Bug in AdvancedIndexing Mixture logprob"),
         ),
         (
             (
@@ -382,7 +381,7 @@ def test_hetero_mixture_binomial(p_val, size):
             (),
             0,
         ),
-        (
+        pytest.param(
             (
                 np.array(0, dtype=pytensor.config.floatX),
                 np.array(1, dtype=pytensor.config.floatX),
@@ -400,6 +399,7 @@ def test_hetero_mixture_binomial(p_val, size):
             (3,),
             (slice(None),),
             1,
+            marks=pytest.mark.xfail(IndexError, reason="Bug in AdvancedIndex Mixture logprob"),
         ),
         (
             (