Fix bug in IfElse Mixture logprob

It did not account for the extra dimension added by the stacking operation, resulting in a logp call to an expanded RV with the original unexpanded value

Author: ricardoV94

pymc/logprob/mixture.py

@@ -443,13 +443,23 @@ def logprob_MixtureRV(
             logp_val = at.set_subtensor(logp_val[idx_m_on_axis], logp_m)
 
     else:
+        # If the stacking operation expands the component RVs, we have
+        # to expand the value and later squeeze the logprob for everything
+        # to work correctly
+        join_axis_val = None if isinstance(join_axis.type, NoneTypeT) else join_axis.data
+
+        if join_axis_val is not None:
+            value = at.expand_dims(value, axis=join_axis_val)
+
         logp_val = 0.0
         for i, comp_rv in enumerate(comp_rvs):
             comp_logp = logprob(comp_rv, value)
+            if join_axis_val is not None:
+                comp_logp = at.squeeze(comp_logp, axis=join_axis_val)
             logp_val += ifelse(
                 at.eq(indices[0], i),
                 comp_logp,
-                at.zeros_like(value),
+                at.zeros_like(comp_logp),
             )
 
     return logp_val

pymc/tests/logprob/test_mixture.py

@@ -226,6 +226,26 @@ def test_hetero_mixture_binomial(p_val, size):
             (),
             0,
         ),
+        # Degenerate vector mixture components, scalar index
+        (
+            (
+                np.array([0], dtype=pytensor.config.floatX),
+                np.array(1, dtype=pytensor.config.floatX),
+            ),
+            (
+                np.array([0.5], dtype=pytensor.config.floatX),
+                np.array(0.5, dtype=pytensor.config.floatX),
+            ),
+            (
+                np.array([100], dtype=pytensor.config.floatX),
+                np.array(1, dtype=pytensor.config.floatX),
+            ),
+            np.array([0.1, 0.5, 0.4], dtype=pytensor.config.floatX),
+            None,
+            (),
+            (),
+            0,
+        ),
         # Scalar mixture components, vector index
         (
             (
@@ -443,16 +463,25 @@ def test_hetero_mixture_categorical(
     gamma_sp = sp.gamma(Y_args[0], scale=1 / Y_args[1])
     norm_2_sp = sp.norm(loc=Z_args[0], scale=Z_args[1])
 
+    # Handle scipy annoying squeeze of random draws
+    real_comp_size = tuple(X_rv.shape.eval())
+
     for i in range(10):
         i_val = CategoricalRV.rng_fn(test_val_rng, p_val, idx_size)
 
         indices_val = list(extra_indices)
         indices_val.insert(join_axis, i_val)
         indices_val = tuple(indices_val)
 
-        x_val = norm_1_sp.rvs(size=comp_size, random_state=test_val_rng)
-        y_val = gamma_sp.rvs(size=comp_size, random_state=test_val_rng)
-        z_val = norm_2_sp.rvs(size=comp_size, random_state=test_val_rng)
+        x_val = np.broadcast_to(
+            norm_1_sp.rvs(size=comp_size, random_state=test_val_rng), real_comp_size
+        )
+        y_val = np.broadcast_to(
+            gamma_sp.rvs(size=comp_size, random_state=test_val_rng), real_comp_size
+        )
+        z_val = np.broadcast_to(
+            norm_2_sp.rvs(size=comp_size, random_state=test_val_rng), real_comp_size
+        )
 
         component_logps = np.stack(
             [norm_1_sp.logpdf(x_val), gamma_sp.logpdf(y_val), norm_2_sp.logpdf(z_val)],