Replace use of broadcastable with shape in MixtureRV

Author: Larry Dong

aeppl/mixture.py

@@ -180,17 +180,17 @@ def rv_pull_down(x: TensorVariable, dont_touch_vars=None) -> TensorVariable:
 class MixtureRV(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_shape")
 
-    def __init__(self, indices_end_idx, out_dtype, out_broadcastable):
+    def __init__(self, indices_end_idx, out_dtype, out_shape):
         super().__init__()
         self.indices_end_idx = indices_end_idx
         self.out_dtype = out_dtype
-        self.out_broadcastable = out_broadcastable
+        self.out_shape = out_shape
 
     def make_node(self, *inputs):
         return Apply(
-            self, list(inputs), [TensorType(self.out_dtype, self.out_broadcastable)()]
+            self, list(inputs), [TensorType(self.out_dtype, self.out_shape)()]
         )
 
     def perform(self, node, inputs, outputs):
@@ -285,7 +285,7 @@ def mixture_replace(fgraph, node):
     mix_op = MixtureRV(
         1 + len(mixing_indices),
         old_mixture_rv.dtype,
-        old_mixture_rv.broadcastable,
+        old_mixture_rv.type.shape,
     )
     new_node = mix_op.make_node(*([join_axis] + mixing_indices + mixture_rvs))
 
@@ -337,7 +337,7 @@ def switch_mixture_replace(fgraph, node):
     mix_op = MixtureRV(
         2,
         old_mixture_rv.dtype,
-        old_mixture_rv.broadcastable,
+        old_mixture_rv.type.shape,
     )
     if node.inputs[0].ndim == 0:
         # as_nontensor_scalar to allow graphs to be identical to mixture sub-graphs
@@ -397,7 +397,7 @@ def ifelse_mixture_replace(fgraph, node):
     mix_op = MixtureRV(
         2,
         old_mixture_rv.dtype,
-        old_mixture_rv.broadcastable,
+        old_mixture_rv.type.shape,
     )
 
     if node.inputs[0].ndim == 0: