# Buggy code
def logsumexp(x, axis=None):
    # Adapted from https://github.com/Theano/Theano/issues/1563
    x_max = tt.max(x, axis=axis, keepdims=True)
    return tt.log(tt.sum(tt.exp(x - x_max), axis=axis, keepdims=True)) + x_max


# Traceback (AttributeError)
---------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)
~/.local/lib/python3.8/site-packages/pymc3/distributions/mixture.py in _comp_modes(self)
289         try:
--> 290             return tt.as_tensor_variable(self.comp_dists.mode)
291         except AttributeError:

AttributeError: 'list' object has no attribute 'mode'

During handling of the above exception, another exception occurred:

TypeError                                 Traceback (most recent call last)
<ipython-input-8-dedf5c958f15> in <module>
8
9     w2 = pm.Dirichlet('w2', np.array([1, 1]))
---> 10     like = pm.Mixture = pm.Mixture('like', w=w2, comp_dists=[mix, a3], observed=np.random.randn(20))

~/.local/lib/python3.8/site-packages/pymc3/distributions/distribution.py in __new__(cls, name, *args, **kwargs)
44                 raise TypeError("observed needs to be data but got: {}".format(type(data)))
45             total_size = kwargs.pop('total_size', None)
---> 46             dist = cls.dist(*args, **kwargs)
47             return model.Var(name, dist, data, total_size)
48         else:

~/.local/lib/python3.8/site-packages/pymc3/distributions/distribution.py in dist(cls, *args, **kwargs)
55     def dist(cls, *args, **kwargs):
56         dist = object.__new__(cls)
---> 57         dist.__init__(*args, **kwargs)
58         return dist
59

~/.local/lib/python3.8/site-packages/pymc3/distributions/mixture.py in __init__(self, w, comp_dists, *args, **kwargs)
139
140         try:
--> 141             comp_modes = self._comp_modes()
142             comp_mode_logps = self.logp(comp_modes)
143             self.mode = comp_modes[tt.argmax(w * comp_mode_logps, axis=-1)]

~/.local/lib/python3.8/site-packages/pymc3/distributions/mixture.py in _comp_modes(self)
290             return tt.as_tensor_variable(self.comp_dists.mode)
291         except AttributeError:
--> 292             return tt.squeeze(tt.stack([comp_dist.mode
293                                         for comp_dist in self.comp_dists],
294                                        axis=-1))

~/.local/lib/python3.8/site-packages/theano/tensor/basic.py in stack(*tensors, **kwargs)
4726         dtype = scal.upcast(*[i.dtype for i in tensors])
4727         return theano.tensor.opt.MakeVector(dtype)(*tensors)
-> 4728     return join(axis, *[shape_padaxis(t, axis) for t in tensors])
4729
4730

~/.local/lib/python3.8/site-packages/theano/tensor/basic.py in join(axis, *tensors_list)
4500         return tensors_list[0]
4501     else:
-> 4502         return join_(axis, *tensors_list)
4503
4504

~/.local/lib/python3.8/site-packages/theano/gof/op.py in __call__(self, *inputs, **kwargs)
613         """
614         return_list = kwargs.pop('return_list', False)
--> 615         node = self.make_node(*inputs, **kwargs)
616
617         if config.compute_test_value != 'off':

~/.local/lib/python3.8/site-packages/theano/tensor/basic.py in make_node(self, *axis_and_tensors)
4232             return tensor(dtype=out_dtype, broadcastable=bcastable)
4233
-> 4234         return self._make_node_internal(
4235             axis, tensors, as_tensor_variable_args, output_maker)
4236

~/.local/lib/python3.8/site-packages/theano/tensor/basic.py in _make_node_internal(self, axis, tensors, as_tensor_variable_args, output_maker)
4299         if not python_all([x.ndim == len(bcastable)
4300                            for x in as_tensor_variable_args[1:]]):
-> 4301             raise TypeError("Join() can only join tensors with the same "
4302                             "number of dimensions.")
4303

TypeError: Join() can only join tensors with the same number of dimensions.

# Fixed code
def logsumexp(x, axis=None, keepdims=True):
    # Adapted from https://github.com/Theano/Theano/issues/1563
    x_max = tt.max(x, axis=axis, keepdims=True)
    res = tt.log(tt.sum(tt.exp(x - x_max), axis=axis, keepdims=True)) + x_max
    return res if keepdims else res.squeeze()