Allow vector-valued indices for switch/ifelse mixture sub-graphs

Author: larryshamalama

aeppl/mixture.py

@@ -337,9 +337,15 @@ def switch_ifelse_mixture_replace(fgraph, node):
         old_mixture_rv.dtype,
         old_mixture_rv.broadcastable,
     )
-    new_node = mix_op.make_node(
-        *([NoneConst, as_nontensor_scalar(node.inputs[0])] + mixture_rvs)
-    )
+
+    if node.inputs[0].ndim == 0:
+        # as_nontensor_scalar to allow graphs to be identical to mixture sub-graphs
+        # created using at.stack and Subtensor indexing
+        new_node = mix_op.make_node(
+            *([NoneConst, as_nontensor_scalar(node.inputs[0])] + mixture_rvs)
+        )
+    else:
+        new_node = mix_op.make_node(*([NoneConst, node.inputs[0]] + mixture_rvs))
 
     new_mixture_rv = new_node.default_output()