Do not raise from internal mixture rewrite

Author: ricardoV94

pymc/logprob/mixture.py

@@ -39,7 +39,7 @@
 import pytensor
 import pytensor.tensor as at
 
-from pytensor.graph.basic import Apply, Variable
+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.rewriting.basic import (
@@ -246,41 +246,29 @@ def get_stack_mixture_vars(
     node: Apply,
 ) -> Tuple[Optional[List[TensorVariable]], Optional[int]]:
     r"""Extract the mixture terms from a `*Subtensor*` applied to stacked `MeasurableVariable`\s."""
-    if not isinstance(node.op, subtensor_ops):
-        return None, None  # pragma: no cover
 
-    join_axis = NoneConst
+    assert isinstance(node.op, subtensor_ops)
+
     joined_rvs = node.inputs[0]
 
     # First, make sure that it's some sort of concatenation
     if not (joined_rvs.owner and isinstance(joined_rvs.owner.op, (MakeVector, Join))):
-        # Node is not a compatible join `Op`
-        return None, join_axis  # pragma: no cover
+        return None, None
 
     if isinstance(joined_rvs.owner.op, MakeVector):
+        join_axis = NoneConst
         mixture_rvs = joined_rvs.owner.inputs
 
     elif isinstance(joined_rvs.owner.op, Join):
-        mixture_rvs = joined_rvs.owner.inputs[1:]
-        join_axis = joined_rvs.owner.inputs[0]
-        try:
-            # TODO: Find better solution to avoid this circular dependency
-            from pymc.pytensorf import constant_fold
-
-            join_axis = int(constant_fold((join_axis,))[0])
-        except ValueError:
-            # TODO: Support symbolic join axes
-            raise NotImplementedError("Symbolic `Join` axes are not supported in mixtures")
+        # TODO: Find better solution to avoid this circular dependency
+        from pymc.pytensorf import constant_fold
 
-        join_axis = at.as_tensor(join_axis)
+        join_axis = joined_rvs.owner.inputs[0]
+        # TODO: Support symbolic join axes. This will raise ValueError if it's not a constant
+        (join_axis,) = constant_fold((join_axis,), raise_not_constant=False)
+        join_axis = at.as_tensor(join_axis, dtype="int64")
 
-    if not all(rv.owner and isinstance(rv.owner.op, MeasurableVariable) for rv in mixture_rvs):
-        # Currently, all mixture components must be `MeasurableVariable` outputs
-        # TODO: Allow constants and make them Dirac-deltas
-        # raise NotImplementedError(
-        #     "All mixture components must be `MeasurableVariable` outputs"
-        # )
-        return None, join_axis
+        mixture_rvs = joined_rvs.owner.inputs[1:]
 
     return mixture_rvs, join_axis
 
@@ -302,33 +290,45 @@ def mixture_replace(fgraph, node):
 
     old_mixture_rv = node.default_output()
 
-    mixture_res, join_axis = get_stack_mixture_vars(node)
+    mixture_rvs, join_axis = get_stack_mixture_vars(node)
+
+    # We don't support symbolic join axis
+    if mixture_rvs is None or not isinstance(join_axis, (NoneTypeT, Constant)):
+        return None
 
-    if mixture_res is None or any(rv in rv_map_feature.rv_values for rv in mixture_res):
+    # Check that all components are MeasurableVariables and none is already conditioned on
+    if not all(
+        (
+            rv.owner is not None
+            and isinstance(rv.owner.op, MeasurableVariable)
+            and rv not in rv_map_feature.rv_values
+        )
+        for rv in mixture_rvs
+    ):
         return None  # pragma: no cover
 
     mixing_indices = node.inputs[1:]
 
     # We loop through mixture components and collect all the array elements
     # that belong to each one (by way of their indices).
-    mixture_rvs = []
-    for i, component_rv in enumerate(mixture_res):
+    new_mixture_rvs = []
+    for i, component_rv in enumerate(mixture_rvs):
 
         # We create custom types for the mixture components and assign them
         # null `get_measurable_outputs` dispatches so that they aren't
         # erroneously encountered in places like `factorized_joint_logprob`.
         new_node = assign_custom_measurable_outputs(component_rv.owner)
         out_idx = component_rv.owner.outputs.index(component_rv)
         new_comp_rv = new_node.outputs[out_idx]
-        mixture_rvs.append(new_comp_rv)
+        new_mixture_rvs.append(new_comp_rv)
 
     # Replace this sub-graph with a `MixtureRV`
     mix_op = MixtureRV(
         1 + len(mixing_indices),
         old_mixture_rv.dtype,
         old_mixture_rv.broadcastable,
     )
-    new_node = mix_op.make_node(*([join_axis] + mixing_indices + mixture_rvs))
+    new_node = mix_op.make_node(*([join_axis] + mixing_indices + new_mixture_rvs))
 
     new_mixture_rv = new_node.default_output()