From 7e9cf7076064d42942fd77cdfdae88cdbb990f0d Mon Sep 17 00:00:00 2001
From: lucianopaz <luciano.paz.neuro@gmail.com>
Date: Thu, 11 Apr 2019 12:31:33 +0200
Subject: [PATCH 01/24] Added shape_utils script

---
 pymc3/distributions/shape_utils.py | 241 +++++++++++++++++++++++++++++
 1 file changed, 241 insertions(+)
 create mode 100644 pymc3/distributions/shape_utils.py

diff --git a/pymc3/distributions/shape_utils.py b/pymc3/distributions/shape_utils.py
new file mode 100644
index 0000000000..c69d6f5935
--- /dev/null
+++ b/pymc3/distributions/shape_utils.py
@@ -0,0 +1,241 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+"""
+A collection of common numpy array shape operations needed for broadcasting
+samples from probability distributions for stochastic nodes in PyMC.
+"""
+import numpy as np
+from .dist_math import to_tuple
+
+
+__all__ = [
+    "broadcast_shapes",
+    "broadcast_distribution_samples",
+    "get_broadcastable_distribution_samples",
+    "broadcast_distribution_samples_shape",
+]
+
+
+def check_shape_type(shape):
+    try:
+        for s in shape:
+            int(s)
+    except Exception:
+        raise TypeError(
+            "Supplied value {} does not represent a valid shape".
+            format(shape)
+        )
+
+
+def broadcast_shapes(*args):
+    """Return the shape resulting from broadcasting multiple shapes.
+    Represents numpy's broadcasting rules.
+    Parameters
+    ----------
+    *args : array-like of int
+        Tuples or arrays or lists representing the shapes of arrays to be
+        broadcast.
+    Returns
+    -------
+    Resulting shape or None if broadcasting is not possible.
+    """
+    x = list(np.atleast_1d(args[0])) if args else ()
+    check_shape_type(x)
+    for arg in args[1:]:
+        y = list(np.atleast_1d(arg))
+        check_shape_type(y)
+        if len(x) < len(y):
+            x, y = y, x
+        x[-len(y):] = [j if i == 1 else i if j == 1 else i if i == j else 0
+                       for i, j in zip(x[-len(y):], y)]
+        if not all(x):
+            return None
+    return tuple(x)
+
+
+def broadcast_distribution_samples(samples, size=None):
+    """Broadcast samples drawn from distributions taking into account the
+    size (i.e. the number of samples) of the draw, which is prepended to
+    the sample's shape.
+
+    Parameters
+    ----------
+    samples: Iterable of ndarrays holding the sampled values
+    size: None, int or tuple (optional)
+        size of the sample set requested.
+
+    Returns
+    -------
+    List of broadcasted sample arrays
+
+    Examples
+    --------
+    .. code-block:: python
+        size = 100
+        sample0 = np.random.randn(size)
+        sample1 = np.random.randn(size, 5)
+        sample2 = np.random.randn(size, 4, 5)
+        out = broadcast_distribution_samples([sample0, sample1, sample2],
+                                             size=size)
+        assert all((o.shape == (size, 4, 5) for o in out))
+        assert np.all(sample0[:, None, None] == out[0])
+        assert np.all(sample1[:, None, :] == out[1])
+        assert np.all(sample2 == out[2])
+
+    .. code-block:: python
+        size = 100
+        sample0 = np.random.randn(size)
+        sample1 = np.random.randn(5)
+        sample2 = np.random.randn(4, 5)
+        out = broadcast_distribution_samples([sample0, sample1, sample2],
+                                             size=size)
+        assert all((o.shape == (size, 4, 5) for o in out))
+        assert np.all(sample0[:, None, None] == out[0])
+        assert np.all(sample1 == out[1])
+        assert np.all(sample2 == out[2])
+    """
+    return np.broadcast_arrays(
+        *get_broadcastable_distribution_samples(samples, size=size)
+    )
+
+
+def get_broadcastable_distribution_samples(samples, size=None):
+    """Get a view of the samples drawn from distributions which adds new axises
+    in between the size prepend and the distribution shape. These views should
+    be able to broadcast the samples from the distrubtions, taking into account
+    the size (i.e. the number of samples) of the draw, which is prepended to
+    the sample's shape.
+
+    Parameters
+    ----------
+    samples: Iterable of ndarrays holding the sampled values
+    size: None, int or tuple (optional)
+        size of the sample set requested.
+
+    Returns
+    -------
+    List of broadcastable views of the sample arrays
+
+    Examples
+    --------
+    .. code-block:: python
+        size = 100
+        sample0 = np.random.randn(size)
+        sample1 = np.random.randn(size, 5)
+        sample2 = np.random.randn(size, 4, 5)
+        out = get_broadcastable_distribution_samples(
+            [sample0, sample1, sample2],
+            size=size
+        )
+        assert np.all(sample0[:, None, None] == out[0])
+        assert np.all(sample1[:, None, :] == out[1])
+        assert np.all(sample2 == out[2])
+
+    .. code-block:: python
+        size = 100
+        sample0 = np.random.randn(size)
+        sample1 = np.random.randn(5)
+        sample2 = np.random.randn(4, 5)
+        out = get_broadcastable_distribution_samples(
+            [sample0, sample1, sample2],
+            size=size
+        )
+        assert np.all(sample0[:, None, None] == out[0])
+        assert np.all(sample1 == out[1])
+        assert np.all(sample2 == out[2])
+    """
+    if size is None:
+        return samples
+    _size = to_tuple(size)
+    # Raw samples shapes
+    p_shapes = [p.shape for p in samples]
+    # samples shapes without the size prepend
+    sp_shapes = [s[len(_size) :] if _size == s[: len(_size)] else s for s in p_shapes]
+    broadcast_shape = np.broadcast(*[np.empty(s) for s in sp_shapes]).shape
+    broadcastable_samples = []
+    for param, p_shape, sp_shape in zip(samples, p_shapes, sp_shapes):
+        if _size == p_shape[: len(_size)]:
+            # If size prepends the shape, then we have to add broadcasting axis
+            # in the middle
+            slicer_head = [slice(None)] * len(_size)
+            slicer_tail = [np.newaxis] * (len(broadcast_shape) - len(sp_shape)) + [
+                slice(None)
+            ] * len(sp_shape)
+        else:
+            # If size does not prepend the shape, then we have leave the
+            # parameter as is
+            slicer_head = []
+            slicer_tail = [slice(None)] * len(sp_shape)
+        broadcastable_samples.append(param[tuple(slicer_head + slicer_tail)])
+    return broadcastable_samples
+
+
+def broadcast_distribution_samples_shape(shapes, size=None):
+    """Get the resulting broadcasted shape for samples drawn from distributions,
+    taking into account the size (i.e. the number of samples) of the draw,
+    which is prepended to the sample's shape.
+
+    Parameters
+    ----------
+    shapes: Iterable of tuples holding the distribution samples shapes
+    size: None, int or tuple (optional)
+        size of the sample set requested.
+
+    Returns
+    -------
+    tuple of the resulting shape
+
+    Examples
+    --------
+    .. code-block:: python
+        size = 100
+        shape0 = (size,)
+        shape1 = (size, 5)
+        shape2 = (size, 4, 5)
+        out = broadcast_distribution_samples_shape([shape0, shape1, shape2],
+                                                   size=size)
+        assert out == (size, 4, 5)
+
+    .. code-block:: python
+        size = 100
+        shape0 = (size,)
+        shape1 = (5,)
+        shape2 = (4, 5)
+        out = broadcast_distribution_samples_shape([shape0, shape1, shape2],
+                                                   size=size)
+        assert out == (size, 4, 5)
+    """
+    if size is None:
+        broadcasted_shape = broadcast_shapes(shapes)
+        if broadcasted_shape is None:
+            raise ValueError(
+                "Cannot broadcast provided shapes {} given size: {}".format(
+                    ", ".join(shapes), size
+                )
+            )
+        return broadcasted_shape
+    _size = to_tuple(size)
+    # samples shapes without the size prepend
+    sp_shapes = [s[len(_size) :] if _size == s[: len(_size)] else s for s in shapes]
+    broadcast_shape = np.broadcast(*[np.empty(s) for s in sp_shapes]).shape
+    broadcastable_shapes = []
+    for shape, sp_shape in zip(shapes, sp_shapes):
+        if _size == shape[: len(_size)]:
+            # If size prepends the shape, then we have to add broadcasting axis
+            # in the middle
+            p_shape = (
+                shape[: len(_size)]
+                + (1,) * (len(broadcast_shape) - len(sp_shape))
+                + shape[len(_size) :]
+            )
+        else:
+            p_shape = shape
+        broadcastable_shapes.append(p_shape)
+    broadcasted_shape = broadcast_shapes(*broadcastable_shapes)
+    if broadcasted_shape is None:
+        raise ValueError(
+            "Cannot broadcast provided shapes {} given size: {}".format(
+                ", ".join(shapes), size
+            )
+        )
+    return broadcasted_shape

From f9dc846d829ec155bc0b6979fd8c3fd6c92ddc4c Mon Sep 17 00:00:00 2001
From: lucianopaz <luciano.paz.neuro@gmail.com>
Date: Thu, 11 Apr 2019 12:32:04 +0200
Subject: [PATCH 02/24] Import from shape_utils and change generate samples

---
 pymc3/distributions/continuous.py   |   4 +-
 pymc3/distributions/discrete.py     |   4 +-
 pymc3/distributions/distribution.py | 113 +++++++---------------------
 pymc3/distributions/mixture.py      |   4 +-
 4 files changed, 35 insertions(+), 90 deletions(-)

diff --git a/pymc3/distributions/continuous.py b/pymc3/distributions/continuous.py
index 0901412b3b..f91a5d0f35 100644
--- a/pymc3/distributions/continuous.py
+++ b/pymc3/distributions/continuous.py
@@ -19,8 +19,8 @@
     alltrue_elemwise, betaln, bound, gammaln, i0e, incomplete_beta, logpow,
     normal_lccdf, normal_lcdf, SplineWrapper, std_cdf, zvalue,
 )
-from .distribution import (Continuous, draw_values, generate_samples,
-                           broadcast_distribution_samples)
+from .distribution import (Continuous, draw_values, generate_samples)
+from .shape_utils import broadcast_distribution_samples
 
 __all__ = ['Uniform', 'Flat', 'HalfFlat', 'Normal', 'TruncatedNormal', 'Beta',
            'Kumaraswamy', 'Exponential', 'Laplace', 'StudentT', 'Cauchy',
diff --git a/pymc3/distributions/discrete.py b/pymc3/distributions/discrete.py
index a6fd0f5f7d..a2a243b3f1 100644
--- a/pymc3/distributions/discrete.py
+++ b/pymc3/distributions/discrete.py
@@ -5,8 +5,8 @@
 
 from pymc3.util import get_variable_name
 from .dist_math import bound, factln, binomln, betaln, logpow, random_choice
-from .distribution import (Discrete, draw_values, generate_samples,
-                           broadcast_distribution_samples)
+from .distribution import Discrete, draw_values, generate_samples
+from .shape_utils import broadcast_distribution_samples
 from pymc3.math import tround, sigmoid, logaddexp, logit, log1pexp
 from ..theanof import floatX, intX
 
diff --git a/pymc3/distributions/distribution.py b/pymc3/distributions/distribution.py
index ff776b0d32..4b665b8ff9 100644
--- a/pymc3/distributions/distribution.py
+++ b/pymc3/distributions/distribution.py
@@ -11,6 +11,12 @@
 )
 from ..vartypes import string_types
 from .dist_math import to_tuple
+from .shape_utils import (
+    get_broadcastable_distribution_samples,
+    broadcast_distribution_samples_shape,
+    broadcast_shapes,
+    broadcast_distribution_samples,
+)
 
 __all__ = ['DensityDist', 'Distribution', 'Continuous', 'Discrete',
            'NoDistribution', 'TensorType', 'draw_values', 'generate_samples']
@@ -286,7 +292,6 @@ def draw_values(params, point=None, size=None):
 
             a) are named parameters in the point
             b) are *RVs with a random method
-
     """
     # Get fast drawable values (i.e. things in point or numbers, arrays,
     # constants or shares, or things that were already drawn in related
@@ -610,31 +615,35 @@ def generate_samples(generator, *args, **kwargs):
     for key in kwargs:
         p = kwargs[key]
         kwargs[key] = p[0] if isinstance(p, tuple) else p
+    size_tup = to_tuple(size)
+    dist_shape = to_tuple(dist_shape)
 
     if broadcast_shape is None:
         inputs = args + tuple(kwargs.values())
-        try:
-            broadcast_shape = np.broadcast(*inputs).shape  # size of generator(size=1)
-        except ValueError:  # sometimes happens if args have shape (500,) and (500, 4)
-            max_dims = max(j.ndim for j in args + tuple(kwargs.values()))
-            args = tuple([j.reshape(j.shape + (1,) * (max_dims - j.ndim)) for j in args])
-            kwargs = {k: v.reshape(v.shape + (1,) * (max_dims - v.ndim)) for k, v in kwargs.items()}
-            inputs = args + tuple(kwargs.values())
-            broadcast_shape = np.broadcast(*inputs).shape  # size of generator(size=1)
+        inputs = get_broadcastable_distribution_samples(
+            inputs + (np.empty(dist_shape),),
+            size=size_tup,
+        )[:-1]
+        broadcast_shape = broadcast_distribution_samples_shape(
+            [inputs[0].shape, dist_shape],
+            size=size_tup
+        )
+        args = inputs[:len(args)]
+        for offset, key in enumerate(kwargs):
+            kwargs[key] = inputs[len(args) + offset]
     # Update kwargs with the keyword arguments that were not broadcasted
     kwargs.update(not_broadcast_kwargs)
 
-    dist_shape = to_tuple(dist_shape)
     broadcast_shape = to_tuple(broadcast_shape)
-    size_tup = to_tuple(size)
 
+    dist_broadcast_shape = broadcast_shapes(dist_shape, broadcast_shape)
     # All inputs are scalars, end up size (size_tup, dist_shape)
     if broadcast_shape in {(), (0,), (1,)}:
         samples = generator(size=size_tup + dist_shape, *args, **kwargs)
     # Inputs already have the right shape. Just get the right size.
-    elif broadcast_shape[-len(dist_shape):] == dist_shape or len(dist_shape) == 0:
-        if size == 1 or (broadcast_shape == size_tup + dist_shape):
-            samples = generator(size=broadcast_shape, *args, **kwargs)
+    elif dist_broadcast_shape or len(dist_shape) == 0:
+        if size == 1:
+            samples = generator(size=dist_broadcast_shape, *args, **kwargs)
         elif dist_shape == broadcast_shape:
             samples = generator(size=size_tup + dist_shape, *args, **kwargs)
         elif len(dist_shape) == 0 and size_tup and broadcast_shape:
@@ -661,8 +670,12 @@ def generate_samples(generator, *args, **kwargs):
                 samples = generator(size=size_tup + broadcast_shape,
                                     *args,
                                     **kwargs)
+        elif dist_broadcast_shape[:len(size_tup)] != size_tup:
+            samples = generator(size=size_tup + dist_broadcast_shape,
+                                *args, **kwargs)
         else:
-            samples = None
+            samples = generator(size=dist_broadcast_shape,
+                                *args, **kwargs)
     # Args have been broadcast correctly, can just ask for the right shape out
     elif dist_shape[-len(broadcast_shape):] == broadcast_shape:
         samples = generator(size=size_tup + dist_shape, *args, **kwargs)
@@ -690,72 +703,4 @@ def generate_samples(generator, *args, **kwargs):
 
     if one_d and samples.shape[-1] == 1:
         samples = samples.reshape(samples.shape[:-1])
-    return np.asarray(samples)
-
-
-def broadcast_distribution_samples(samples, size=None):
-    """Broadcast samples drawn from distributions taking into account the
-    size (i.e. the number of samples) of the draw, which is prepended to
-    the sample's shape.
-
-    Parameters
-    ----------
-    samples: Iterable of ndarrays holding the sampled values
-    size: None, int or tuple (optional)
-        size of the sample set requested.
-
-    Returns
-    -------
-    List of broadcasted sample arrays
-
-    Examples
-    --------
-    .. code-block:: python
-        size = 100
-        sample0 = np.random.randn(size)
-        sample1 = np.random.randn(size, 5)
-        sample2 = np.random.randn(size, 4, 5)
-        out = broadcast_distribution_samples([sample0, sample1, sample2],
-                                             size=size)
-        assert all((o.shape == (size, 4, 5) for o in out))
-        assert np.all(sample0[:, None, None] == out[0])
-        assert np.all(sample1[:, None, :] == out[1])
-        assert np.all(sample2 == out[2])
-
-    .. code-block:: python
-        size = 100
-        sample0 = np.random.randn(size)
-        sample1 = np.random.randn(5)
-        sample2 = np.random.randn(4, 5)
-        out = broadcast_distribution_samples([sample0, sample1, sample2],
-                                             size=size)
-        assert all((o.shape == (size, 4, 5) for o in out))
-        assert np.all(sample0[:, None, None] == out[0])
-        assert np.all(sample1 == out[1])
-        assert np.all(sample2 == out[2])
-    """
-    if size is None:
-        return np.broadcast_arrays(*samples)
-    _size = to_tuple(size)
-    # Raw samples shapes
-    p_shapes = [p.shape for p in samples]
-    # samples shapes without the size prepend
-    sp_shapes = [s[len(_size):] if _size == s[:len(_size)] else s
-                 for s in p_shapes]
-    broadcast_shape = np.broadcast(*[np.empty(s) for s in sp_shapes]).shape
-    broadcasted_samples = []
-    for param, p_shape, sp_shape in zip(samples, p_shapes, sp_shapes):
-        if _size == p_shape[:len(_size)]:
-            # If size prepends the shape, then we have to add broadcasting axis
-            # in the middle
-            slicer_head = [slice(None)] * len(_size)
-            slicer_tail = ([np.newaxis] * (len(broadcast_shape) -
-                                           len(sp_shape)) +
-                           [slice(None)] * len(sp_shape))
-        else:
-            # If size does not prepend the shape, then we have leave the
-            # parameter as is
-            slicer_head = []
-            slicer_tail = [slice(None)] * len(sp_shape)
-        broadcasted_samples.append(param[tuple(slicer_head + slicer_tail)])
-    return np.broadcast_arrays(*broadcasted_samples)
+    return np.asarray(samples)
\ No newline at end of file
diff --git a/pymc3/distributions/mixture.py b/pymc3/distributions/mixture.py
index ee18fed8b7..4bae5aeb82 100644
--- a/pymc3/distributions/mixture.py
+++ b/pymc3/distributions/mixture.py
@@ -8,8 +8,8 @@
 from .dist_math import bound, random_choice, to_tuple
 from .distribution import (Discrete, Distribution, draw_values,
                            generate_samples, _DrawValuesContext,
-                           _DrawValuesContextBlocker,
-                           broadcast_distribution_samples)
+                           _DrawValuesContextBlocker)
+from .shape_utils import broadcast_distribution_samples
 from .continuous import get_tau_sigma, Normal
 from ..theanof import _conversion_map
 

From 6fababc0779be5701359dfba552226636672870a Mon Sep 17 00:00:00 2001
From: lucianopaz <luciano.paz.neuro@gmail.com>
Date: Thu, 11 Apr 2019 13:12:03 +0200
Subject: [PATCH 03/24] Fixed lint and errors

---
 pymc3/distributions/distribution.py |  6 +++---
 pymc3/distributions/shape_utils.py  | 28 +++++++++++++++++++++-------
 2 files changed, 24 insertions(+), 10 deletions(-)

diff --git a/pymc3/distributions/distribution.py b/pymc3/distributions/distribution.py
index 4b665b8ff9..2da91e3806 100644
--- a/pymc3/distributions/distribution.py
+++ b/pymc3/distributions/distribution.py
@@ -15,7 +15,6 @@
     get_broadcastable_distribution_samples,
     broadcast_distribution_samples_shape,
     broadcast_shapes,
-    broadcast_distribution_samples,
 )
 
 __all__ = ['DensityDist', 'Distribution', 'Continuous', 'Discrete',
@@ -625,10 +624,10 @@ def generate_samples(generator, *args, **kwargs):
             size=size_tup,
         )[:-1]
         broadcast_shape = broadcast_distribution_samples_shape(
-            [inputs[0].shape, dist_shape],
+            [i.shape for i in inputs] + [dist_shape],
             size=size_tup
         )
-        args = inputs[:len(args)]
+        args = tuple(inputs[:len(args)])
         for offset, key in enumerate(kwargs):
             kwargs[key] = inputs[len(args) + offset]
     # Update kwargs with the keyword arguments that were not broadcasted
@@ -639,6 +638,7 @@ def generate_samples(generator, *args, **kwargs):
     dist_broadcast_shape = broadcast_shapes(dist_shape, broadcast_shape)
     # All inputs are scalars, end up size (size_tup, dist_shape)
     if broadcast_shape in {(), (0,), (1,)}:
+        print(size_tup, dist_shape, [a.shape for a in args + tuple(kwargs.values())])
         samples = generator(size=size_tup + dist_shape, *args, **kwargs)
     # Inputs already have the right shape. Just get the right size.
     elif dist_broadcast_shape or len(dist_shape) == 0:
diff --git a/pymc3/distributions/shape_utils.py b/pymc3/distributions/shape_utils.py
index c69d6f5935..c4a78c8663 100644
--- a/pymc3/distributions/shape_utils.py
+++ b/pymc3/distributions/shape_utils.py
@@ -27,7 +27,7 @@ def check_shape_type(shape):
         )
 
 
-def broadcast_shapes(*args):
+def broadcast_shapes(*args, raise_exception=False):
     """Return the shape resulting from broadcasting multiple shapes.
     Represents numpy's broadcasting rules.
     Parameters
@@ -49,7 +49,13 @@ def broadcast_shapes(*args):
         x[-len(y):] = [j if i == 1 else i if j == 1 else i if i == j else 0
                        for i, j in zip(x[-len(y):], y)]
         if not all(x):
-            return None
+            if raise_exception:
+                raise ValueError(
+                    "Supplied shapes {} do not broadcast together".
+                    format(", ".join(["{}".format(a) for a in args]))
+                )
+            else:
+                return None
     return tuple(x)
 
 
@@ -144,6 +150,7 @@ def get_broadcastable_distribution_samples(samples, size=None):
         assert np.all(sample1 == out[1])
         assert np.all(sample2 == out[2])
     """
+    samples = [np.atleast_1d(p) for p in samples]
     if size is None:
         return samples
     _size = to_tuple(size)
@@ -151,7 +158,7 @@ def get_broadcastable_distribution_samples(samples, size=None):
     p_shapes = [p.shape for p in samples]
     # samples shapes without the size prepend
     sp_shapes = [s[len(_size) :] if _size == s[: len(_size)] else s for s in p_shapes]
-    broadcast_shape = np.broadcast(*[np.empty(s) for s in sp_shapes]).shape
+    broadcast_shape = broadcast_shapes(*sp_shapes, raise_exception=True)
     broadcastable_samples = []
     for param, p_shape, sp_shape in zip(samples, p_shapes, sp_shapes):
         if _size == p_shape[: len(_size)]:
@@ -206,18 +213,25 @@ def broadcast_distribution_samples_shape(shapes, size=None):
         assert out == (size, 4, 5)
     """
     if size is None:
-        broadcasted_shape = broadcast_shapes(shapes)
+        broadcasted_shape = broadcast_shapes(*shapes)
         if broadcasted_shape is None:
             raise ValueError(
                 "Cannot broadcast provided shapes {} given size: {}".format(
-                    ", ".join(shapes), size
+                    ", ".join(["{}".format(s) for s in shapes]), size
                 )
             )
         return broadcasted_shape
     _size = to_tuple(size)
     # samples shapes without the size prepend
     sp_shapes = [s[len(_size) :] if _size == s[: len(_size)] else s for s in shapes]
-    broadcast_shape = np.broadcast(*[np.empty(s) for s in sp_shapes]).shape
+    try:
+        broadcast_shape = broadcast_shapes(*sp_shapes, raise_exception=True)
+    except ValueError:
+        raise ValueError(
+            "Cannot broadcast provided shapes {} given size: {}".format(
+                ", ".join(["{}".format(s) for s in shapes]), size
+            )
+        )
     broadcastable_shapes = []
     for shape, sp_shape in zip(shapes, sp_shapes):
         if _size == shape[: len(_size)]:
@@ -235,7 +249,7 @@ def broadcast_distribution_samples_shape(shapes, size=None):
     if broadcasted_shape is None:
         raise ValueError(
             "Cannot broadcast provided shapes {} given size: {}".format(
-                ", ".join(shapes), size
+                ", ".join(["{}".format(s) for s in shapes]), size
             )
         )
     return broadcasted_shape

From 041046a3d4e5e280151fdf0b999f67361cdef2a1 Mon Sep 17 00:00:00 2001
From: lucianopaz <luciano.paz.neuro@gmail.com>
Date: Thu, 11 Apr 2019 13:19:06 +0200
Subject: [PATCH 04/24] Still errors with broadcast shapes

---
 pymc3/distributions/distribution.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/pymc3/distributions/distribution.py b/pymc3/distributions/distribution.py
index 2da91e3806..a5aaa27471 100644
--- a/pymc3/distributions/distribution.py
+++ b/pymc3/distributions/distribution.py
@@ -622,9 +622,9 @@ def generate_samples(generator, *args, **kwargs):
         inputs = get_broadcastable_distribution_samples(
             inputs + (np.empty(dist_shape),),
             size=size_tup,
-        )[:-1]
+        )
         broadcast_shape = broadcast_distribution_samples_shape(
-            [i.shape for i in inputs] + [dist_shape],
+            [i.shape for i in inputs],
             size=size_tup
         )
         args = tuple(inputs[:len(args)])
@@ -638,7 +638,7 @@ def generate_samples(generator, *args, **kwargs):
     dist_broadcast_shape = broadcast_shapes(dist_shape, broadcast_shape)
     # All inputs are scalars, end up size (size_tup, dist_shape)
     if broadcast_shape in {(), (0,), (1,)}:
-        print(size_tup, dist_shape, [a.shape for a in args + tuple(kwargs.values())])
+        print(size_tup, dist_shape, broadcast_shape, [a.shape for a in args + tuple(kwargs.values())])
         samples = generator(size=size_tup + dist_shape, *args, **kwargs)
     # Inputs already have the right shape. Just get the right size.
     elif dist_broadcast_shape or len(dist_shape) == 0:

From 6215aa4439c648bae4d63f47366fba4acc29e824 Mon Sep 17 00:00:00 2001
From: lucianopaz <luciano.paz.neuro@gmail.com>
Date: Fri, 12 Apr 2019 11:51:10 +0200
Subject: [PATCH 05/24] Fixed triangular errors

---
 pymc3/distributions/continuous.py   |  2 +-
 pymc3/distributions/distribution.py | 14 +++++++-------
 pymc3/distributions/shape_utils.py  | 25 ++++++++++++++-----------
 3 files changed, 22 insertions(+), 19 deletions(-)

diff --git a/pymc3/distributions/continuous.py b/pymc3/distributions/continuous.py
index f91a5d0f35..abfda57044 100644
--- a/pymc3/distributions/continuous.py
+++ b/pymc3/distributions/continuous.py
@@ -3369,7 +3369,7 @@ def random(self, point=None, size=None):
         scale = upper - lower
         c_ = (c - lower) / scale
         return generate_samples(stats.triang.rvs, c=c_, loc=lower, scale=scale,
-                                size=size, dist_shape=self.shape, random_state=None)
+                                size=size, dist_shape=self.shape)
 
     def logp(self, value):
         """
diff --git a/pymc3/distributions/distribution.py b/pymc3/distributions/distribution.py
index a5aaa27471..d108bfb950 100644
--- a/pymc3/distributions/distribution.py
+++ b/pymc3/distributions/distribution.py
@@ -14,7 +14,7 @@
 from .shape_utils import (
     get_broadcastable_distribution_samples,
     broadcast_distribution_samples_shape,
-    broadcast_shapes,
+    shapes_broadcasting,
 )
 
 __all__ = ['DensityDist', 'Distribution', 'Continuous', 'Discrete',
@@ -606,8 +606,8 @@ def generate_samples(generator, *args, **kwargs):
     not_broadcast_kwargs = kwargs.pop('not_broadcast_kwargs', None)
     if not_broadcast_kwargs is None:
         not_broadcast_kwargs = dict()
-    if size is None:
-        size = 1
+#    if size is None:
+#        size = 1
 
     args = tuple(p[0] if isinstance(p, tuple) else p for p in args)
 
@@ -635,14 +635,13 @@ def generate_samples(generator, *args, **kwargs):
 
     broadcast_shape = to_tuple(broadcast_shape)
 
-    dist_broadcast_shape = broadcast_shapes(dist_shape, broadcast_shape)
+    dist_broadcast_shape = shapes_broadcasting(dist_shape, broadcast_shape)
     # All inputs are scalars, end up size (size_tup, dist_shape)
     if broadcast_shape in {(), (0,), (1,)}:
-        print(size_tup, dist_shape, broadcast_shape, [a.shape for a in args + tuple(kwargs.values())])
         samples = generator(size=size_tup + dist_shape, *args, **kwargs)
     # Inputs already have the right shape. Just get the right size.
-    elif dist_broadcast_shape or len(dist_shape) == 0:
-        if size == 1:
+    elif dist_broadcast_shape is not None or len(dist_shape) == 0:
+        if size is None:
             samples = generator(size=dist_broadcast_shape, *args, **kwargs)
         elif dist_shape == broadcast_shape:
             samples = generator(size=size_tup + dist_shape, *args, **kwargs)
@@ -697,6 +696,7 @@ def generate_samples(generator, *args, **kwargs):
         '''.format(size=size, size_tup=size_tup, dist_shape=dist_shape, broadcast_shape=broadcast_shape, test=broadcast_shape[:len(size_tup)] == size_tup))
 
     # reshape samples here
+    samples = np.atleast_1d(samples)
     if samples.shape[0] == 1 and size == 1:
         if len(samples.shape) > len(dist_shape) and samples.shape[-len(dist_shape):] == dist_shape:
             samples = samples.reshape(samples.shape[1:])
diff --git a/pymc3/distributions/shape_utils.py b/pymc3/distributions/shape_utils.py
index c4a78c8663..8d56ba857f 100644
--- a/pymc3/distributions/shape_utils.py
+++ b/pymc3/distributions/shape_utils.py
@@ -9,7 +9,7 @@
 
 
 __all__ = [
-    "broadcast_shapes",
+    "shapes_broadcasting",
     "broadcast_distribution_samples",
     "get_broadcastable_distribution_samples",
     "broadcast_distribution_samples_shape",
@@ -17,17 +17,19 @@
 
 
 def check_shape_type(shape):
+    out = []
     try:
         for s in shape:
-            int(s)
+            out.append(int(s))
     except Exception:
         raise TypeError(
             "Supplied value {} does not represent a valid shape".
             format(shape)
         )
+    return tuple(out)
 
 
-def broadcast_shapes(*args, raise_exception=False):
+def shapes_broadcasting(*args, raise_exception=False):
     """Return the shape resulting from broadcasting multiple shapes.
     Represents numpy's broadcasting rules.
     Parameters
@@ -40,14 +42,15 @@ def broadcast_shapes(*args, raise_exception=False):
     Resulting shape or None if broadcasting is not possible.
     """
     x = list(np.atleast_1d(args[0])) if args else ()
-    check_shape_type(x)
+    x = list(check_shape_type(x))
     for arg in args[1:]:
         y = list(np.atleast_1d(arg))
-        check_shape_type(y)
+        y = list(check_shape_type(y))
         if len(x) < len(y):
             x, y = y, x
-        x[-len(y):] = [j if i == 1 else i if j == 1 else i if i == j else 0
-                       for i, j in zip(x[-len(y):], y)]
+        if len(y) > 0:
+            x[-len(y):] = [j if i == 1 else i if j == 1 else i if i == j else 0
+                           for i, j in zip(x[-len(y):], y)]
         if not all(x):
             if raise_exception:
                 raise ValueError(
@@ -158,7 +161,7 @@ def get_broadcastable_distribution_samples(samples, size=None):
     p_shapes = [p.shape for p in samples]
     # samples shapes without the size prepend
     sp_shapes = [s[len(_size) :] if _size == s[: len(_size)] else s for s in p_shapes]
-    broadcast_shape = broadcast_shapes(*sp_shapes, raise_exception=True)
+    broadcast_shape = shapes_broadcasting(*sp_shapes, raise_exception=True)
     broadcastable_samples = []
     for param, p_shape, sp_shape in zip(samples, p_shapes, sp_shapes):
         if _size == p_shape[: len(_size)]:
@@ -213,7 +216,7 @@ def broadcast_distribution_samples_shape(shapes, size=None):
         assert out == (size, 4, 5)
     """
     if size is None:
-        broadcasted_shape = broadcast_shapes(*shapes)
+        broadcasted_shape = shapes_broadcasting(*shapes)
         if broadcasted_shape is None:
             raise ValueError(
                 "Cannot broadcast provided shapes {} given size: {}".format(
@@ -225,7 +228,7 @@ def broadcast_distribution_samples_shape(shapes, size=None):
     # samples shapes without the size prepend
     sp_shapes = [s[len(_size) :] if _size == s[: len(_size)] else s for s in shapes]
     try:
-        broadcast_shape = broadcast_shapes(*sp_shapes, raise_exception=True)
+        broadcast_shape = shapes_broadcasting(*sp_shapes, raise_exception=True)
     except ValueError:
         raise ValueError(
             "Cannot broadcast provided shapes {} given size: {}".format(
@@ -245,7 +248,7 @@ def broadcast_distribution_samples_shape(shapes, size=None):
         else:
             p_shape = shape
         broadcastable_shapes.append(p_shape)
-    broadcasted_shape = broadcast_shapes(*broadcastable_shapes)
+    broadcasted_shape = shapes_broadcasting(*broadcastable_shapes)
     if broadcasted_shape is None:
         raise ValueError(
             "Cannot broadcast provided shapes {} given size: {}".format(

From a581f636d9772bbab4c32e0925002ddfa4211065 Mon Sep 17 00:00:00 2001
From: lucianopaz <luciano.paz.neuro@gmail.com>
Date: Tue, 16 Apr 2019 17:27:25 +0200
Subject: [PATCH 06/24] Fixed distrubutions random errors.

---
 pymc3/distributions/distribution.py | 117 +++++------
 pymc3/distributions/shape_utils.py  | 301 ++++++++++++++++++----------
 2 files changed, 240 insertions(+), 178 deletions(-)

diff --git a/pymc3/distributions/distribution.py b/pymc3/distributions/distribution.py
index d108bfb950..d223fb1a19 100644
--- a/pymc3/distributions/distribution.py
+++ b/pymc3/distributions/distribution.py
@@ -12,9 +12,8 @@
 from ..vartypes import string_types
 from .dist_math import to_tuple
 from .shape_utils import (
-    get_broadcastable_distribution_samples,
-    broadcast_distribution_samples_shape,
-    shapes_broadcasting,
+    get_broadcastable_dist_samples,
+    broadcast_dist_samples_shape,
 )
 
 __all__ = ['DensityDist', 'Distribution', 'Continuous', 'Discrete',
@@ -602,105 +601,89 @@ def generate_samples(generator, *args, **kwargs):
     dist_shape = kwargs.pop('dist_shape', ())
     one_d = _is_one_d(dist_shape)
     size = kwargs.pop('size', None)
+#    if size is None:
+#        size = 1
     broadcast_shape = kwargs.pop('broadcast_shape', None)
     not_broadcast_kwargs = kwargs.pop('not_broadcast_kwargs', None)
     if not_broadcast_kwargs is None:
         not_broadcast_kwargs = dict()
-#    if size is None:
-#        size = 1
 
+    # Parse out raw input parameters for the generator
     args = tuple(p[0] if isinstance(p, tuple) else p for p in args)
-
     for key in kwargs:
         p = kwargs[key]
         kwargs[key] = p[0] if isinstance(p, tuple) else p
+
+    # Convert size and dist_shape to tuples
     size_tup = to_tuple(size)
     dist_shape = to_tuple(dist_shape)
 
     if broadcast_shape is None:
+        # If broadcast_shape is not explicitly provided, it is inferred as the
+        # broadcasted shape of the input parameter and dist_shape, taking into
+        # account the potential size prefix
         inputs = args + tuple(kwargs.values())
-        inputs = get_broadcastable_distribution_samples(
-            inputs + (np.empty(dist_shape),),
-            size=size_tup,
-        )
-        broadcast_shape = broadcast_distribution_samples_shape(
-            [i.shape for i in inputs],
+        broadcast_shape = broadcast_dist_samples_shape(
+            [np.atleast_1d(i).shape for i in inputs] + [dist_shape],
             size=size_tup
         )
+        # We do this instead of broadcast_distribution_samples to avoid
+        # creating a dummy array with dist_shape in memory
+        inputs = get_broadcastable_dist_samples(
+            inputs,
+            size=size_tup,
+            must_bcast_with=broadcast_shape,
+        )
+        # We modify the arguments with their broadcasted counterparts
         args = tuple(inputs[:len(args)])
         for offset, key in enumerate(kwargs):
             kwargs[key] = inputs[len(args) + offset]
     # Update kwargs with the keyword arguments that were not broadcasted
     kwargs.update(not_broadcast_kwargs)
 
+    # We ensure that broadcast_shape is a tuple
     broadcast_shape = to_tuple(broadcast_shape)
 
-    dist_broadcast_shape = shapes_broadcasting(dist_shape, broadcast_shape)
-    # All inputs are scalars, end up size (size_tup, dist_shape)
-    if broadcast_shape in {(), (0,), (1,)}:
-        samples = generator(size=size_tup + dist_shape, *args, **kwargs)
-    # Inputs already have the right shape. Just get the right size.
-    elif dist_broadcast_shape is not None or len(dist_shape) == 0:
-        if size is None:
-            samples = generator(size=dist_broadcast_shape, *args, **kwargs)
-        elif dist_shape == broadcast_shape:
-            samples = generator(size=size_tup + dist_shape, *args, **kwargs)
-        elif len(dist_shape) == 0 and size_tup and broadcast_shape:
-            # There is no dist_shape (scalar distribution) but the parameters
-            # broadcast shape and size_tup determine the size to provide to
-            # the generator
-            if broadcast_shape[:len(size_tup)] == size_tup:
-                # Input's dist_shape is scalar, but it has size repetitions.
-                # So now the size matches but we have to manually broadcast to
-                # the right dist_shape
-                samples = [generator(*args, **kwargs)]
-                if samples[0].shape == broadcast_shape:
-                    samples = samples[0]
-                else:
-                    suffix = broadcast_shape[len(size_tup):] + dist_shape
-                    samples.extend([generator(*args, **kwargs).
-                                    reshape(broadcast_shape)[..., np.newaxis]
-                                    for _ in range(np.prod(suffix,
-                                                           dtype=int) - 1)])
-                    samples = np.hstack(samples).reshape(size_tup + suffix)
-            else:
-                # The parameter shape is given, but we have to concatenate it
-                # with the size tuple
-                samples = generator(size=size_tup + broadcast_shape,
-                                    *args,
-                                    **kwargs)
-        elif dist_broadcast_shape[:len(size_tup)] != size_tup:
-            samples = generator(size=size_tup + dist_broadcast_shape,
-                                *args, **kwargs)
-        else:
-            samples = generator(size=dist_broadcast_shape,
-                                *args, **kwargs)
-    # Args have been broadcast correctly, can just ask for the right shape out
-    elif dist_shape[-len(broadcast_shape):] == broadcast_shape:
-        samples = generator(size=size_tup + dist_shape, *args, **kwargs)
-    # Inputs have the right size, have to manually broadcast to the right dist_shape
-    elif broadcast_shape[:len(size_tup)] == size_tup:
-        suffix = broadcast_shape[len(size_tup):] + dist_shape
-        samples = [generator(*args, **kwargs).reshape(size_tup + (1,)) for _ in range(np.prod(suffix, dtype=int))]
-        samples = np.hstack(samples).reshape(size_tup + suffix)
+    if dist_shape[:len(size_tup)] == size_tup:
+        # dist_shape is prepended with size_tup. This is not a consequence
+        # of the parameters being drawn size_tup times! By chance, the
+        # distribution's shape has its first elements equal to size_tup.
+        # This means that we must prepend the size_tup to dist_shape, and
+        # check if that broadcasts well with the parameters
+        _dist_shape = size_tup + dist_shape
     else:
-        samples = None
-
-    if samples is None:
+        _dist_shape = dist_shape
+    try:
+        dist_bcast_shape = broadcast_dist_samples_shape(
+            [_dist_shape, broadcast_shape],
+            size=size,
+        )
+    except (ValueError, TypeError):
         raise TypeError('''Attempted to generate values with incompatible shapes:
             size: {size}
             size_tup: {size_tup}
             broadcast_shape[:len(size_tup)] == size_tup: {test}
             dist_shape: {dist_shape}
             broadcast_shape: {broadcast_shape}
-        '''.format(size=size, size_tup=size_tup, dist_shape=dist_shape, broadcast_shape=broadcast_shape, test=broadcast_shape[:len(size_tup)] == size_tup))
+        '''.format(size=size,
+                   size_tup=size_tup,
+                   dist_shape=dist_shape,
+                   broadcast_shape=broadcast_shape,
+                   test=broadcast_shape[:len(size_tup)] == size_tup)
+        )
+    if dist_bcast_shape[:len(size_tup)] == size_tup:
+        samples = generator(size=dist_bcast_shape, *args, **kwargs)
+    else:
+        samples = generator(size=size_tup + dist_bcast_shape, *args, **kwargs)
+    samples = np.asarray(samples)
 
     # reshape samples here
-    samples = np.atleast_1d(samples)
     if samples.shape[0] == 1 and size == 1:
-        if len(samples.shape) > len(dist_shape) and samples.shape[-len(dist_shape):] == dist_shape:
+        if (len(samples.shape) > len(dist_shape) and
+            samples.shape[-len(dist_shape):] == dist_shape[-len(dist_shape):]
+        ):
             samples = samples.reshape(samples.shape[1:])
 
-    if one_d and samples.shape[-1] == 1:
+    if one_d and samples.ndim > 0 and samples.shape[-1] == 1:
         samples = samples.reshape(samples.shape[:-1])
     return np.asarray(samples)
\ No newline at end of file
diff --git a/pymc3/distributions/shape_utils.py b/pymc3/distributions/shape_utils.py
index 8d56ba857f..7549b54e47 100644
--- a/pymc3/distributions/shape_utils.py
+++ b/pymc3/distributions/shape_utils.py
@@ -10,21 +10,21 @@
 
 __all__ = [
     "shapes_broadcasting",
+    "broadcast_dist_samples_shape",
+    "get_broadcastable_dist_samples",
     "broadcast_distribution_samples",
-    "get_broadcastable_distribution_samples",
-    "broadcast_distribution_samples_shape",
+    "broadcast_dist_samples_to",
 ]
 
 
-def check_shape_type(shape):
+def _check_shape_type(shape):
     out = []
     try:
         for s in shape:
             out.append(int(s))
     except Exception:
         raise TypeError(
-            "Supplied value {} does not represent a valid shape".
-            format(shape)
+            "Supplied value {} does not represent a valid shape".format(shape)
         )
     return tuple(out)
 
@@ -42,123 +42,177 @@ def shapes_broadcasting(*args, raise_exception=False):
     Resulting shape or None if broadcasting is not possible.
     """
     x = list(np.atleast_1d(args[0])) if args else ()
-    x = list(check_shape_type(x))
+    x = list(_check_shape_type(x))
     for arg in args[1:]:
         y = list(np.atleast_1d(arg))
-        y = list(check_shape_type(y))
+        y = list(_check_shape_type(y))
         if len(x) < len(y):
             x, y = y, x
         if len(y) > 0:
-            x[-len(y):] = [j if i == 1 else i if j == 1 else i if i == j else 0
-                           for i, j in zip(x[-len(y):], y)]
+            x[-len(y) :] = [
+                j if i == 1 else i if j == 1 else i if i == j else 0
+                for i, j in zip(x[-len(y) :], y)
+            ]
         if not all(x):
             if raise_exception:
                 raise ValueError(
-                    "Supplied shapes {} do not broadcast together".
-                    format(", ".join(["{}".format(a) for a in args]))
+                    "Supplied shapes {} do not broadcast together".format(
+                        ", ".join(["{}".format(a) for a in args])
+                    )
                 )
             else:
                 return None
     return tuple(x)
 
 
-def broadcast_distribution_samples(samples, size=None):
-    """Broadcast samples drawn from distributions taking into account the
-    size (i.e. the number of samples) of the draw, which is prepended to
-    the sample's shape.
+def broadcast_dist_samples_shape(shapes, size=None):
+    """Get the resulting broadcasted shape for samples drawn from distributions,
+    taking into account the size (i.e. the number of samples) of the draw,
+    which is prepended to the sample's shape.
 
     Parameters
     ----------
-    samples: Iterable of ndarrays holding the sampled values
+    shapes: Iterable of tuples holding the distribution samples shapes
     size: None, int or tuple (optional)
         size of the sample set requested.
 
     Returns
     -------
-    List of broadcasted sample arrays
+    tuple of the resulting shape
 
     Examples
     --------
     .. code-block:: python
         size = 100
-        sample0 = np.random.randn(size)
-        sample1 = np.random.randn(size, 5)
-        sample2 = np.random.randn(size, 4, 5)
-        out = broadcast_distribution_samples([sample0, sample1, sample2],
-                                             size=size)
-        assert all((o.shape == (size, 4, 5) for o in out))
-        assert np.all(sample0[:, None, None] == out[0])
-        assert np.all(sample1[:, None, :] == out[1])
-        assert np.all(sample2 == out[2])
+        shape0 = (size,)
+        shape1 = (size, 5)
+        shape2 = (size, 4, 5)
+        out = broadcast_dist_samples_shape([shape0, shape1, shape2],
+                                                   size=size)
+        assert out == (size, 4, 5)
 
     .. code-block:: python
         size = 100
-        sample0 = np.random.randn(size)
-        sample1 = np.random.randn(5)
-        sample2 = np.random.randn(4, 5)
-        out = broadcast_distribution_samples([sample0, sample1, sample2],
-                                             size=size)
-        assert all((o.shape == (size, 4, 5) for o in out))
-        assert np.all(sample0[:, None, None] == out[0])
-        assert np.all(sample1 == out[1])
-        assert np.all(sample2 == out[2])
+        shape0 = (size,)
+        shape1 = (5,)
+        shape2 = (4, 5)
+        out = broadcast_dist_samples_shape([shape0, shape1, shape2],
+                                                   size=size)
+        assert out == (size, 4, 5)
     """
-    return np.broadcast_arrays(
-        *get_broadcastable_distribution_samples(samples, size=size)
-    )
+    if size is None:
+        broadcasted_shape = shapes_broadcasting(*shapes)
+        if broadcasted_shape is None:
+            raise ValueError(
+                "Cannot broadcast provided shapes {} given size: {}".format(
+                    ", ".join(["{}".format(s) for s in shapes]), size
+                )
+            )
+        return broadcasted_shape
+    shapes = [_check_shape_type(s) for s in shapes]
+    _size = to_tuple(size)
+    # samples shapes without the size prepend
+    sp_shapes = [s[len(_size) :] if _size == s[: len(_size)] else s for s in shapes]
+    try:
+        broadcast_shape = shapes_broadcasting(*sp_shapes, raise_exception=True)
+    except ValueError:
+        raise ValueError(
+            "Cannot broadcast provided shapes {} given size: {}".format(
+                ", ".join(["{}".format(s) for s in shapes]), size
+            )
+        )
+    broadcastable_shapes = []
+    for shape, sp_shape in zip(shapes, sp_shapes):
+        if _size == shape[: len(_size)]:
+            # If size prepends the shape, then we have to add broadcasting axis
+            # in the middle
+            p_shape = (
+                shape[: len(_size)]
+                + (1,) * (len(broadcast_shape) - len(sp_shape))
+                + shape[len(_size) :]
+            )
+        else:
+            p_shape = shape
+        broadcastable_shapes.append(p_shape)
+    broadcasted_shape = shapes_broadcasting(*broadcastable_shapes)
+    if broadcasted_shape is None:
+        raise ValueError(
+            "Cannot broadcast provided shapes {} given size: {}".format(
+                ", ".join(["{}".format(s) for s in shapes]), size
+            )
+        )
+    return broadcasted_shape
 
 
-def get_broadcastable_distribution_samples(samples, size=None):
+def get_broadcastable_dist_samples(
+    samples, size=None, must_bcast_with=None, return_out_shape=False
+):
     """Get a view of the samples drawn from distributions which adds new axises
     in between the size prepend and the distribution shape. These views should
-    be able to broadcast the samples from the distrubtions, taking into account
+    be able to broadcast the samples from the distrubtions taking into account
     the size (i.e. the number of samples) of the draw, which is prepended to
-    the sample's shape.
+    the sample's shape. Optionally, one can supply an extra `must_bcast_with`
+    to try to force samples to be able to broadcast with a given shape.
 
     Parameters
     ----------
     samples: Iterable of ndarrays holding the sampled values
     size: None, int or tuple (optional)
         size of the sample set requested.
+    must_bcast_with: None, int or tuple (optional)
+        Tuple shape to which the samples must be able to broadcast
+    return_out_shape: bool (optional)
+        If `True`, this function also returns the output's shape and not only
+        samples views.
 
     Returns
     -------
-    List of broadcastable views of the sample arrays
+    broadcastable_samples: List of the broadcasted sample arrays
+    broadcast_shape: If `return_out_shape` is `True`, the resulting broadcast
+        shape is returned.
 
     Examples
     --------
     .. code-block:: python
+        must_bcast_with = (3, 1, 5)
         size = 100
         sample0 = np.random.randn(size)
         sample1 = np.random.randn(size, 5)
         sample2 = np.random.randn(size, 4, 5)
-        out = get_broadcastable_distribution_samples(
+        out = broadcast_dist_samples_to(
             [sample0, sample1, sample2],
-            size=size
+            size=size,
+            must_bcast_with=must_bcast_with,
         )
-        assert np.all(sample0[:, None, None] == out[0])
-        assert np.all(sample1[:, None, :] == out[1])
-        assert np.all(sample2 == out[2])
+        assert np.all((o.shape == (size, 3, 4, 5) for o in out))
+        assert np.all(sample0[:, None, None, None] == out[0])
+        assert np.all(sample1[:, None, None] == out[1])
+        assert np.all(sample2[:, None] == out[2])
 
     .. code-block:: python
         size = 100
+        must_bcast_with = (3, 1, 5)
         sample0 = np.random.randn(size)
         sample1 = np.random.randn(5)
         sample2 = np.random.randn(4, 5)
-        out = get_broadcastable_distribution_samples(
+        out = broadcast_dist_samples_to(
             [sample0, sample1, sample2],
-            size=size
+            size=size,
+            must_bcast_with=must_bcast_with,
         )
-        assert np.all(sample0[:, None, None] == out[0])
+        assert np.all((o.shape == (size, 3, 4, 5) for o in out))
+        assert np.all(sample0[:, None, None, None] == out[0])
         assert np.all(sample1 == out[1])
         assert np.all(sample2 == out[2])
     """
-    samples = [np.atleast_1d(p) for p in samples]
-    if size is None:
+    samples = [np.asarray(p) for p in samples]
+    if size is None and must_bcast_with is None:
         return samples
     _size = to_tuple(size)
+    must_bcast_with = to_tuple(must_bcast_with)
     # Raw samples shapes
-    p_shapes = [p.shape for p in samples]
+    p_shapes = [p.shape for p in samples] + [_check_shape_type(must_bcast_with)]
+    out_shape = broadcast_dist_samples_shape(p_shapes, size=size)
     # samples shapes without the size prepend
     sp_shapes = [s[len(_size) :] if _size == s[: len(_size)] else s for s in p_shapes]
     broadcast_shape = shapes_broadcasting(*sp_shapes, raise_exception=True)
@@ -177,82 +231,107 @@ def get_broadcastable_distribution_samples(samples, size=None):
             slicer_head = []
             slicer_tail = [slice(None)] * len(sp_shape)
         broadcastable_samples.append(param[tuple(slicer_head + slicer_tail)])
-    return broadcastable_samples
+    if return_out_shape:
+        return broadcastable_samples, out_shape
+    else:
+        return broadcastable_samples
 
 
-def broadcast_distribution_samples_shape(shapes, size=None):
-    """Get the resulting broadcasted shape for samples drawn from distributions,
-    taking into account the size (i.e. the number of samples) of the draw,
-    which is prepended to the sample's shape.
+def broadcast_distribution_samples(samples, size=None):
+    """Broadcast samples drawn from distributions taking into account the
+    size (i.e. the number of samples) of the draw, which is prepended to
+    the sample's shape.
 
     Parameters
     ----------
-    shapes: Iterable of tuples holding the distribution samples shapes
+    samples: Iterable of ndarrays holding the sampled values
     size: None, int or tuple (optional)
         size of the sample set requested.
 
     Returns
     -------
-    tuple of the resulting shape
+    List of broadcasted sample arrays
 
     Examples
     --------
     .. code-block:: python
         size = 100
-        shape0 = (size,)
-        shape1 = (size, 5)
-        shape2 = (size, 4, 5)
-        out = broadcast_distribution_samples_shape([shape0, shape1, shape2],
-                                                   size=size)
-        assert out == (size, 4, 5)
+        sample0 = np.random.randn(size)
+        sample1 = np.random.randn(size, 5)
+        sample2 = np.random.randn(size, 4, 5)
+        out = broadcast_distribution_samples([sample0, sample1, sample2],
+                                             size=size)
+        assert all((o.shape == (size, 4, 5) for o in out))
+        assert np.all(sample0[:, None, None] == out[0])
+        assert np.all(sample1[:, None, :] == out[1])
+        assert np.all(sample2 == out[2])
 
     .. code-block:: python
         size = 100
-        shape0 = (size,)
-        shape1 = (5,)
-        shape2 = (4, 5)
-        out = broadcast_distribution_samples_shape([shape0, shape1, shape2],
-                                                   size=size)
-        assert out == (size, 4, 5)
+        sample0 = np.random.randn(size)
+        sample1 = np.random.randn(5)
+        sample2 = np.random.randn(4, 5)
+        out = broadcast_distribution_samples([sample0, sample1, sample2],
+                                             size=size)
+        assert all((o.shape == (size, 4, 5) for o in out))
+        assert np.all(sample0[:, None, None] == out[0])
+        assert np.all(sample1 == out[1])
+        assert np.all(sample2 == out[2])
     """
-    if size is None:
-        broadcasted_shape = shapes_broadcasting(*shapes)
-        if broadcasted_shape is None:
-            raise ValueError(
-                "Cannot broadcast provided shapes {} given size: {}".format(
-                    ", ".join(["{}".format(s) for s in shapes]), size
-                )
-            )
-        return broadcasted_shape
-    _size = to_tuple(size)
-    # samples shapes without the size prepend
-    sp_shapes = [s[len(_size) :] if _size == s[: len(_size)] else s for s in shapes]
-    try:
-        broadcast_shape = shapes_broadcasting(*sp_shapes, raise_exception=True)
-    except ValueError:
-        raise ValueError(
-            "Cannot broadcast provided shapes {} given size: {}".format(
-                ", ".join(["{}".format(s) for s in shapes]), size
-            )
+    return np.broadcast_arrays(*get_broadcastable_dist_samples(samples, size=size))
+
+
+def broadcast_dist_samples_to(to_shape, samples, size=None):
+    """Broadcast samples drawn from distributions to a given shape, taking into
+    account the size (i.e. the number of samples) of the draw, which is
+    prepended to the sample's shape.
+
+    Parameters
+    ----------
+    to_shape: Tuple shape to which to broadcast the samples
+    samples: Iterable of ndarrays holding the sampled values
+    size: None, int or tuple (optional)
+        size of the sample set requested.
+
+    Returns
+    -------
+    List of the broadcasted sample arrays
+
+    Examples
+    --------
+    .. code-block:: python
+        to_shape = (3, 1, 5)
+        size = 100
+        sample0 = np.random.randn(size)
+        sample1 = np.random.randn(size, 5)
+        sample2 = np.random.randn(size, 4, 5)
+        out = broadcast_dist_samples_to(
+            to_shape,
+            [sample0, sample1, sample2],
+            size=size
         )
-    broadcastable_shapes = []
-    for shape, sp_shape in zip(shapes, sp_shapes):
-        if _size == shape[: len(_size)]:
-            # If size prepends the shape, then we have to add broadcasting axis
-            # in the middle
-            p_shape = (
-                shape[: len(_size)]
-                + (1,) * (len(broadcast_shape) - len(sp_shape))
-                + shape[len(_size) :]
-            )
-        else:
-            p_shape = shape
-        broadcastable_shapes.append(p_shape)
-    broadcasted_shape = shapes_broadcasting(*broadcastable_shapes)
-    if broadcasted_shape is None:
-        raise ValueError(
-            "Cannot broadcast provided shapes {} given size: {}".format(
-                ", ".join(["{}".format(s) for s in shapes]), size
-            )
+        assert np.all((o.shape == (size, 3, 4, 5) for o in out))
+        assert np.all(sample0[:, None, None, None] == out[0])
+        assert np.all(sample1[:, None, None] == out[1])
+        assert np.all(sample2[:, None] == out[2])
+
+    .. code-block:: python
+        size = 100
+        to_shape = (3, 1, 5)
+        sample0 = np.random.randn(size)
+        sample1 = np.random.randn(5)
+        sample2 = np.random.randn(4, 5)
+        out = broadcast_dist_samples_to(
+            to_shape,
+            [sample0, sample1, sample2],
+            size=size
         )
-    return broadcasted_shape
+        assert np.all((o.shape == (size, 3, 4, 5) for o in out))
+        assert np.all(sample0[:, None, None, None] == out[0])
+        assert np.all(sample1 == out[1])
+        assert np.all(sample2 == out[2])
+    """
+    samples, to_shape = get_broadcastable_dist_samples(
+        samples, size=size, must_bcast_with=to_shape, return_out_shape=True
+    )
+    return [np.broadcast_to(o, to_shape) for o in samples]

From a46ceaaa08dd8c04519d42436d983d57e334a41a Mon Sep 17 00:00:00 2001
From: lucianopaz <luciano.paz.neuro@gmail.com>
Date: Tue, 16 Apr 2019 18:56:28 +0200
Subject: [PATCH 07/24] Fixed Multinomial error

---
 pymc3/distributions/distribution.py |   2 -
 pymc3/distributions/multivariate.py | 124 +++++++++-------------------
 pymc3/tests/test_sampling.py        |  15 ++--
 3 files changed, 45 insertions(+), 96 deletions(-)

diff --git a/pymc3/distributions/distribution.py b/pymc3/distributions/distribution.py
index d223fb1a19..4118ad4777 100644
--- a/pymc3/distributions/distribution.py
+++ b/pymc3/distributions/distribution.py
@@ -601,8 +601,6 @@ def generate_samples(generator, *args, **kwargs):
     dist_shape = kwargs.pop('dist_shape', ())
     one_d = _is_one_d(dist_shape)
     size = kwargs.pop('size', None)
-#    if size is None:
-#        size = 1
     broadcast_shape = kwargs.pop('broadcast_shape', None)
     not_broadcast_kwargs = kwargs.pop('not_broadcast_kwargs', None)
     if not_broadcast_kwargs is None:
diff --git a/pymc3/distributions/multivariate.py b/pymc3/distributions/multivariate.py
index d773036d77..9a7c9b2b5d 100755
--- a/pymc3/distributions/multivariate.py
+++ b/pymc3/distributions/multivariate.py
@@ -21,7 +21,7 @@
 from ..model import Deterministic
 from .continuous import ChiSquared, Normal
 from .special import gammaln, multigammaln
-from .dist_math import bound, logpow, factln
+from .dist_math import bound, logpow, factln, to_tuple
 from ..math import kron_dot, kron_diag, kron_solve_lower, kronecker
 
 
@@ -543,104 +543,54 @@ def __init__(self, n, p, *args, **kwargs):
                                 diff[inc_bool_arr.nonzero()])
         self.mode = mode
 
-    def _random(self, n, p, size=None):
+    def _random(self, n, p, size=None, raw_size=None):
         original_dtype = p.dtype
         # Set float type to float64 for numpy. This change is related to numpy issue #8317 (https://github.com/numpy/numpy/issues/8317)
         p = p.astype('float64')
         # Now, re-normalize all of the values in float64 precision. This is done inside the conditionals
+        p /= np.sum(p, axis=-1, keepdims=True)
 
-        # np.random.multinomial needs `n` to be a scalar int and `p` a
-        # sequence
-        if p.ndim == 1 and (n.ndim == 0 or (n.ndim == 1 and n.shape[0] == 1)):
-            # If `n` is already a scalar and `p` is a sequence, then just
-            # return np.multinomial with some size handling
-            p = p / p.sum()
-            if size is not None:
-                if size == p.shape:
-                    size = None
-                elif size[-len(p.shape):] == p.shape:
-                    size = size[:len(size) - len(p.shape)]
-            randnum = np.random.multinomial(n, p, size=size)
-            return randnum.astype(original_dtype)
-        # The shapes of `p` and `n` must be broadcasted by hand depending on
-        # their ndim. We will assume that the last axis of the `p` array will
-        # be the sequence to feed into np.random.multinomial. The other axis
-        # will only have to be iterated over.
-        if n.ndim == p.ndim:
-            # p and n have the same ndim, so n.shape[-1] must be 1
-            if n.shape[-1] != 1:
-                raise ValueError('If n and p have the same number of '
-                                 'dimensions, the last axis of n must be '
-                                 'have len 1. Got {} instead.\n'
-                                 'n.shape = {}\n'
-                                 'p.shape = {}.'.format(n.shape[-1],
-                                                        n.shape,
-                                                        p.shape))
-            n_p_shape = np.broadcast(np.empty(p.shape[:-1]),
-                                     np.empty(n.shape[:-1])).shape
-            p = np.broadcast_to(p, n_p_shape + (p.shape[-1],))
-            n = np.broadcast_to(n, n_p_shape + (1,))
-        elif n.ndim == p.ndim - 1:
-            # n has the number of dimensions of p for the iteration, these must
-            # broadcast together
-            n_p_shape = np.broadcast(np.empty(p.shape[:-1]),
-                                     n).shape
-            p = np.broadcast_to(p, n_p_shape + (p.shape[-1],))
-            n = np.broadcast_to(n, n_p_shape + (1,))
-        elif p.ndim == 1:
-            # p only has the sequence array. We extend it with the dimensions
-            # of n
-            n_p_shape = n.shape
-            p = np.broadcast_to(p, n_p_shape + (p.shape[-1],))
-            n = np.broadcast_to(n, n_p_shape + (1,))
-        elif n.ndim == 0 or (n.dim == 1 and n.shape[0] == 1):
-            # n is a scalar. We extend it with the dimensions of p
-            n_p_shape = p.shape[:-1]
-            n = np.broadcast_to(n, n_p_shape + (1,))
-        else:
-            # There is no clear rule to broadcast p and n so we raise an error
-            raise ValueError('Incompatible shapes of n and p.\n'
-                             'n.shape = {}\n'
-                             'p.shape = {}'.format(n.shape, p.shape))
-
-        # Check what happens with size
-        if size is not None:
-            if size == p.shape:
-                size = None
-                _size = 1
-            elif size[-len(p.shape):] == p.shape:
-                size = size[:len(size) - len(p.shape)]
-                _size = np.prod(size)
-            else:
-                _size = np.prod(size)
-        else:
-            _size = 1
+        # Thanks to the default shape handling done in generate_values, the last
+        # axis of n is a dummy axis that allows it to broadcast well with p
+        n = np.broadcast_to(n, size)
+        p = np.broadcast_to(p, size)
+        n = n[..., 0]
 
-        # We now flatten p and n up to the last dimension
-        p_shape = p.shape
-        p = np.reshape(p, (np.prod(n_p_shape), -1))
-        n = np.reshape(n, (np.prod(n_p_shape), -1))
-        # We renormalize p
-        p = p / p.sum(axis=1, keepdims=True)
-        # We iterate calls to np.random.multinomial
-        randnum = np.asarray([
-            np.random.multinomial(nn, pp, size=_size)
-            for (nn, pp) in zip(n, p)
-        ])
-        # We swap the iteration axis with the _size axis
-        randnum = np.moveaxis(randnum, 1, 0)
-        # We reshape the random numbers to the corresponding size + p_shape
-        if size is None:
-            randnum = np.reshape(randnum, p_shape)
+        # np.random.multinomial needs `n` to be a scalar int and `p` a
+        # sequence so we semi flatten them and iterate over them
+        size_ = to_tuple(raw_size)
+        if p.ndim > len(size_) and p.shape[:len(size_)] == size_:
+            # p and n have the size_ prepend so we don't need it in np.random
+            n_ = n.reshape([-1])
+            p_ = p.reshape([-1, p.shape[-1]])
+            samples = np.array([
+                np.random.multinomial(nn, pp)
+                for nn, pp in zip(n_, p_)
+            ])
+            samples = samples.reshape(p.shape)
         else:
-            randnum = np.reshape(randnum, size + p_shape)
+            # p and n don't have the size prepend
+            n_ = n.reshape([-1])
+            p_ = p.reshape([-1, p.shape[-1]])
+            samples = np.array([
+                np.random.multinomial(nn, pp, size=size_)
+                for nn, pp in zip(n_, p_)
+            ])
+            samples = np.moveaxis(samples, 0, -1)
+            samples = samples.reshape(size + p.shape)
         # We cast back to the original dtype
-        return randnum.astype(original_dtype)
+        return samples.astype(original_dtype)
 
     def random(self, point=None, size=None):
         n, p = draw_values([self.n, self.p], point=point, size=size)
-        samples = generate_samples(self._random, n, p,
+        n = np.asarray(n)
+        p = np.asarray(p)
+        # n should broadcast with p[..., 0] (the last axis in p is not shared
+        # with n for broadcasting). To do this, we add a dummy axis at the end
+        # of n
+        samples = generate_samples(self._random, n[..., None], p,
                                    dist_shape=self.shape,
+                                   not_broadcast_kwargs={'raw_size': size},
                                    size=size)
         return samples
 
diff --git a/pymc3/tests/test_sampling.py b/pymc3/tests/test_sampling.py
index d47a8c8abd..f55b3247f6 100644
--- a/pymc3/tests/test_sampling.py
+++ b/pymc3/tests/test_sampling.py
@@ -486,13 +486,14 @@ def test_multivariate2(self):
             probs = pm.Dirichlet("probs", a=np.ones(6), shape=6)
             obs = pm.Multinomial("obs", n=100, p=probs, observed=mn_data)
             burned_trace = pm.sample(20, tune=10, cores=1)
-        sim_priors = pm.sample_prior_predictive(samples=20, model=dm_model)
-        sim_ppc = pm.sample_posterior_predictive(
-            burned_trace, samples=20, model=dm_model
-        )
-        assert sim_priors["probs"].shape == (20, 6)
-        assert sim_priors["obs"].shape == (20, 6)
-        assert sim_ppc["obs"].shape == (20,) + obs.distribution.shape
+        sim_priors = pm.sample_prior_predictive(samples=20,
+                                                model=dm_model)
+        sim_ppc = pm.sample_posterior_predictive(burned_trace,
+                                                 samples=20,
+                                                 model=dm_model)
+        assert sim_priors['probs'].shape == (20, 6)
+        assert sim_priors['obs'].shape == (20,) + obs.distribution.shape
+        assert sim_ppc['obs'].shape == (20,) + obs.distribution.shape
 
     def test_layers(self):
         with pm.Model() as model:

From fffb89c2ce1a8da7c0db95f4713d3efb13e5862c Mon Sep 17 00:00:00 2001
From: lucianopaz <luciano.paz.neuro@gmail.com>
Date: Wed, 17 Apr 2019 09:56:59 +0200
Subject: [PATCH 08/24] Fixed multinomial n shape error

---
 pymc3/distributions/distribution.py |  2 +-
 pymc3/distributions/multivariate.py | 12 +-----------
 2 files changed, 2 insertions(+), 12 deletions(-)

diff --git a/pymc3/distributions/distribution.py b/pymc3/distributions/distribution.py
index 4118ad4777..c3a32e90f2 100644
--- a/pymc3/distributions/distribution.py
+++ b/pymc3/distributions/distribution.py
@@ -622,7 +622,7 @@ def generate_samples(generator, *args, **kwargs):
         # account the potential size prefix
         inputs = args + tuple(kwargs.values())
         broadcast_shape = broadcast_dist_samples_shape(
-            [np.atleast_1d(i).shape for i in inputs] + [dist_shape],
+            [np.asarray(i).shape for i in inputs] + [dist_shape],
             size=size_tup
         )
         # We do this instead of broadcast_distribution_samples to avoid
diff --git a/pymc3/distributions/multivariate.py b/pymc3/distributions/multivariate.py
index 9a7c9b2b5d..cefea4a7e1 100755
--- a/pymc3/distributions/multivariate.py
+++ b/pymc3/distributions/multivariate.py
@@ -520,11 +520,6 @@ def __init__(self, n, p, *args, **kwargs):
         n = np.squeeze(n) # works also if n is a tensor
 
         if len(self.shape) > 1:
-            m = self.shape[-2]
-            # try:
-            #     assert n.shape == (m,)
-            # except (AttributeError, AssertionError):
-            #     n = n * tt.ones(m)
             self.n = tt.shape_padright(n)
             self.p = p if p.ndim > 1 else tt.shape_padleft(p)
         elif n.ndim == 1:
@@ -583,12 +578,7 @@ def _random(self, n, p, size=None, raw_size=None):
 
     def random(self, point=None, size=None):
         n, p = draw_values([self.n, self.p], point=point, size=size)
-        n = np.asarray(n)
-        p = np.asarray(p)
-        # n should broadcast with p[..., 0] (the last axis in p is not shared
-        # with n for broadcasting). To do this, we add a dummy axis at the end
-        # of n
-        samples = generate_samples(self._random, n[..., None], p,
+        samples = generate_samples(self._random, n, p,
                                    dist_shape=self.shape,
                                    not_broadcast_kwargs={'raw_size': size},
                                    size=size)

From 4d452b6a2683b9fc0d0b5d7b184137b7ce02453d Mon Sep 17 00:00:00 2001
From: lucianopaz <luciano.paz.neuro@gmail.com>
Date: Wed, 24 Apr 2019 18:36:01 +0200
Subject: [PATCH 09/24] Added shape broadcasting tests

---
 RELEASE-NOTES.md                   |   3 +
 pymc3/distributions/shape_utils.py |  14 ++--
 pymc3/tests/test_shape_handling.py | 102 +++++++++++++++++++++++++++++
 3 files changed, 114 insertions(+), 5 deletions(-)
 create mode 100644 pymc3/tests/test_shape_handling.py

diff --git a/RELEASE-NOTES.md b/RELEASE-NOTES.md
index 2a2afc01bc..83e3540197 100644
--- a/RELEASE-NOTES.md
+++ b/RELEASE-NOTES.md
@@ -8,6 +8,7 @@
 - Add function `set_data` to update variables defined as `Data`.
 - `Mixture` now supports mixtures of multidimensional probability distributions, not just lists of 1D distributions.
 - `GLM.from_formula` and `LinearComponent.from_formula` can extract variables from the calling scope. Customizable via the new `eval_env` argument. Fixing #3382.
+- Added the `distributions.shape_utils` module with functions used to help broadcast samples drawn from distributions using the `size` keyword argument.
 
 ### Maintenance
 - All occurances of `sd` as a parameter name have been renamed to `sigma`. `sd` will continue to function for backwards compatibility.
@@ -31,6 +32,8 @@
 - Add `sigma`, `tau`, and `sd` to signature of `NormalMixture`.
 - Resolved issue #3248. Set default lower and upper values of -inf and inf for pm.distributions.continuous.TruncatedNormal. This avoids errors caused by their previous values of None.
 - Resolved issue #3399. Converted all calls to `pm.distributions.bound._ContinuousBounded` and `pm.distributions.bound._DiscreteBounded` to use only and all positional arguments.
+- Restructured `distributions.distribution.generate_samples` to use the `shape_utils` module. This solves issues #3421 and #3147 by using the `size` aware broadcating functions in `shape_utils`.
+- Fixed the `Multinomial.random` and `Multinomial.random_` methods to make them compatible with the new `generate_samples` function. In the process, a bug of the `Multinomial.random_` shape handling was discovered and fixed.
 
 ### Deprecations
 
diff --git a/pymc3/distributions/shape_utils.py b/pymc3/distributions/shape_utils.py
index 7549b54e47..c1e6056876 100644
--- a/pymc3/distributions/shape_utils.py
+++ b/pymc3/distributions/shape_utils.py
@@ -20,8 +20,14 @@
 def _check_shape_type(shape):
     out = []
     try:
+        shape = np.atleast_1d(shape)
         for s in shape:
-            out.append(int(s))
+            if isinstance(s, np.ndarray) and s.ndim > 0:
+                raise TypeError("Value {} is not a valid integer".format(s))
+            o = int(s)
+            if o != s:
+                raise TypeError("Value {} is not a valid integer".format(s))
+            out.append(o)
     except Exception:
         raise TypeError(
             "Supplied value {} does not represent a valid shape".format(shape)
@@ -41,11 +47,9 @@ def shapes_broadcasting(*args, raise_exception=False):
     -------
     Resulting shape or None if broadcasting is not possible.
     """
-    x = list(np.atleast_1d(args[0])) if args else ()
-    x = list(_check_shape_type(x))
+    x = list(_check_shape_type(args[0])) if args else ()
     for arg in args[1:]:
-        y = list(np.atleast_1d(arg))
-        y = list(_check_shape_type(y))
+        y = list(_check_shape_type(arg))
         if len(x) < len(y):
             x, y = y, x
         if len(y) > 0:
diff --git a/pymc3/tests/test_shape_handling.py b/pymc3/tests/test_shape_handling.py
new file mode 100644
index 0000000000..c607d7ae5f
--- /dev/null
+++ b/pymc3/tests/test_shape_handling.py
@@ -0,0 +1,102 @@
+import pytest
+import itertools
+import numpy as np
+import pymc3 as pm
+from pymc3.distributions.shape_utils import (
+    shapes_broadcasting,
+    broadcast_dist_samples_shape,
+    get_broadcastable_dist_samples,
+    broadcast_distribution_samples,
+    broadcast_dist_samples_to,
+)
+from pymc3.distributions.dist_math import to_tuple
+
+test_shapes = [
+    (tuple(), (1,), (4,), (5, 4)),
+    (tuple(), (1,), (7,), (5, 4)),
+    (tuple(), (1,), (1, 4), (5, 4)),
+    (tuple(), (1,), (10, 4), (5, 4)),
+    (tuple(), (1,), (10,), (5, 4)),
+    (tuple(), (1,), (1, 4), (5, 4)),
+    (tuple(), (1,), (3, 4), (5, 4)),
+    (tuple(), (1,), (1, 1, 4), (5, 4)),
+    (tuple(), (1,), (10, 1, 4), (5, 4)),
+    (tuple(), (1,), (5, 1), (5, 4)),
+    (tuple(), (1,), (5, 3), (5, 4)),
+    (tuple(), (1,), (10, 5, 1), (5, 4)),
+    (tuple(), (1,), (10, 5, 3), (5, 4)),
+    (tuple(), (1,), (10, 5, 4), (5, 4)),
+    (tuple(), (1,), (5, 4), (5, 4)),
+]
+test_sizes = [
+    None,
+    tuple(),
+    1,
+    (1,),
+    10,
+    (10,),
+    (1, 1),
+    (10, 1),
+    (7,),
+    (1, 10),
+    (5,),
+    (5, 4),
+]
+
+
+class TestShapesBroadcasting:
+    @pytest.mark.parametrize(
+        "bad_input",
+        [None, [None], "asd", 3.6, {1: 2}, {3}, [8, [8]], "3", ["3"], np.array([[2]])],
+        ids=str,
+    )
+    def test_type_check_raises(self, bad_input):
+        with pytest.raises(TypeError):
+            shapes_broadcasting(bad_input, tuple(), raise_exception=True)
+            shapes_broadcasting(bad_input, tuple(), raise_exception=False)
+
+    @pytest.mark.parametrize("raise_exception", [False, True], ids=str)
+    def test_type_check_success(self, raise_exception):
+        inputs = [3, 3.0, tuple(), [3], (3,), np.array(3), np.array([3])]
+        out = shapes_broadcasting(*inputs, raise_exception=raise_exception)
+        assert out == (3,)
+
+    @pytest.mark.parametrize(
+        ["shapes", "raise_exception"],
+        itertools.product(test_shapes, [False, True]),
+        ids=str,
+    )
+    def test_broadcasting(self, shapes, raise_exception):
+        try:
+            expected_out = np.broadcast(*[np.empty(s) for s in shapes]).shape
+        except ValueError:
+            expected_out = None
+        if expected_out is None:
+            if raise_exception:
+                with pytest.raises(ValueError):
+                    shapes_broadcasting(*shapes, raise_exception=raise_exception)
+            else:
+                out = shapes_broadcasting(*shapes, raise_exception=raise_exception)
+                assert out is None
+        else:
+            out = shapes_broadcasting(*shapes, raise_exception=raise_exception)
+            assert out == expected_out
+
+    @pytest.mark.parametrize(
+        ["size", "shapes"], itertools.product(test_sizes, test_shapes), ids=str
+    )
+    def test_broadcast_dist_samples_shape(self, size, shapes):
+        size_ = to_tuple(size)
+        shapes_ = [s if s[:len(size_)] != size_ else s[len(size_):] for s in shapes]
+        try:
+            expected_out = np.broadcast(*[np.empty(s) for s in shapes_]).shape
+        except ValueError:
+            expected_out = None
+        if expected_out is not None and any((s[: len(size_)] == size_ for s in shapes)):
+            expected_out = size_ + expected_out
+        if expected_out is None:
+            with pytest.raises(ValueError):
+                broadcast_dist_samples_shape(shapes, size=size)
+        else:
+            out = broadcast_dist_samples_shape(shapes, size=size)
+            assert out == expected_out

From e2029181fc5d33d071fdb21d0e2287a620dbb2f6 Mon Sep 17 00:00:00 2001
From: lucianopaz <luciano.paz.neuro@gmail.com>
Date: Thu, 25 Apr 2019 18:15:20 +0200
Subject: [PATCH 10/24] Added broadcast samples tests

---
 pymc3/distributions/shape_utils.py |  8 ++--
 pymc3/tests/test_shape_handling.py | 75 +++++++++++++++++++++++++++++-
 2 files changed, 76 insertions(+), 7 deletions(-)

diff --git a/pymc3/distributions/shape_utils.py b/pymc3/distributions/shape_utils.py
index c1e6056876..47b5da3098 100644
--- a/pymc3/distributions/shape_utils.py
+++ b/pymc3/distributions/shape_utils.py
@@ -116,7 +116,7 @@ def broadcast_dist_samples_shape(shapes, size=None):
     shapes = [_check_shape_type(s) for s in shapes]
     _size = to_tuple(size)
     # samples shapes without the size prepend
-    sp_shapes = [s[len(_size) :] if _size == s[: len(_size)] else s for s in shapes]
+    sp_shapes = [s[len(_size) :] if _size == s[: min([len(_size), len(s)])] else s for s in shapes]
     try:
         broadcast_shape = shapes_broadcasting(*sp_shapes, raise_exception=True)
     except ValueError:
@@ -210,19 +210,17 @@ def get_broadcastable_dist_samples(
         assert np.all(sample2 == out[2])
     """
     samples = [np.asarray(p) for p in samples]
-    if size is None and must_bcast_with is None:
-        return samples
     _size = to_tuple(size)
     must_bcast_with = to_tuple(must_bcast_with)
     # Raw samples shapes
     p_shapes = [p.shape for p in samples] + [_check_shape_type(must_bcast_with)]
     out_shape = broadcast_dist_samples_shape(p_shapes, size=size)
     # samples shapes without the size prepend
-    sp_shapes = [s[len(_size) :] if _size == s[: len(_size)] else s for s in p_shapes]
+    sp_shapes = [s[len(_size) :] if _size == s[: min([len(_size), len(s)])] else s for s in p_shapes]
     broadcast_shape = shapes_broadcasting(*sp_shapes, raise_exception=True)
     broadcastable_samples = []
     for param, p_shape, sp_shape in zip(samples, p_shapes, sp_shapes):
-        if _size == p_shape[: len(_size)]:
+        if _size == p_shape[: min([len(_size), len(p_shape)])]:
             # If size prepends the shape, then we have to add broadcasting axis
             # in the middle
             slicer_head = [slice(None)] * len(_size)
diff --git a/pymc3/tests/test_shape_handling.py b/pymc3/tests/test_shape_handling.py
index c607d7ae5f..7dd6e6fa0e 100644
--- a/pymc3/tests/test_shape_handling.py
+++ b/pymc3/tests/test_shape_handling.py
@@ -87,12 +87,17 @@ def test_broadcasting(self, shapes, raise_exception):
     )
     def test_broadcast_dist_samples_shape(self, size, shapes):
         size_ = to_tuple(size)
-        shapes_ = [s if s[:len(size_)] != size_ else s[len(size_):] for s in shapes]
+        shapes_ = [
+            s if s[: min([len(size_), len(s)])] != size_ else s[len(size_) :]
+            for s in shapes
+        ]
         try:
             expected_out = np.broadcast(*[np.empty(s) for s in shapes_]).shape
         except ValueError:
             expected_out = None
-        if expected_out is not None and any((s[: len(size_)] == size_ for s in shapes)):
+        if expected_out is not None and any(
+            (s[: min([len(size_), len(s)])] == size_ for s in shapes)
+        ):
             expected_out = size_ + expected_out
         if expected_out is None:
             with pytest.raises(ValueError):
@@ -100,3 +105,69 @@ def test_broadcast_dist_samples_shape(self, size, shapes):
         else:
             out = broadcast_dist_samples_shape(shapes, size=size)
             assert out == expected_out
+
+
+@pytest.fixture(
+    scope="module", params=itertools.product(test_sizes, test_shapes), ids=str
+)
+def samples_to_broadcast(request):
+    size, shapes = request.param
+    samples = [np.empty(s) for s in shapes]
+    try:
+        broadcast_shape = broadcast_dist_samples_shape(shapes, size=size)
+    except ValueError:
+        broadcast_shape = None
+    return size, samples, broadcast_shape
+
+
+@pytest.fixture(
+    scope="module", params=(tuple(), (1,), (10, 5, 4), (10, 1, 1, 5, 1)), ids=str
+)
+def samples_to_broadcast_to(request, samples_to_broadcast):
+    to_shape = request.param
+    size, samples, broadcast_shape = samples_to_broadcast
+    if broadcast_shape is not None:
+        try:
+            broadcast_shape = broadcast_dist_samples_shape(
+                [broadcast_shape, to_shape], size=size
+            )
+        except ValueError:
+            broadcast_shape = None
+    return to_shape, size, samples, broadcast_shape
+
+
+class TestSamplesBroadcasting:
+    def test_broadcast_distribution_samples(self, samples_to_broadcast):
+        size, samples, broadcast_shape = samples_to_broadcast
+        if broadcast_shape is not None:
+            outs = broadcast_distribution_samples(samples, size=size)
+            assert all((o.shape == broadcast_shape for o in outs))
+        else:
+            with pytest.raises(ValueError):
+                broadcast_distribution_samples(samples, size=size)
+
+    def test_get_broadcastable_dist_samples(self, samples_to_broadcast):
+        size, samples, broadcast_shape = samples_to_broadcast
+        if broadcast_shape is not None:
+            outs, out_shape = get_broadcastable_dist_samples(
+                samples, size=size, return_out_shape=True
+            )
+            assert out_shape == broadcast_shape
+            for o in outs:
+                for oshape, bshape in itertools.zip_longest(
+                    reversed(o.shape), reversed(broadcast_shape), fillvalue=1
+                ):
+                    assert oshape == bshape or oshape == 1
+            assert shapes_broadcasting(*[o.shape for o in outs]) == broadcast_shape
+        else:
+            with pytest.raises(ValueError):
+                get_broadcastable_dist_samples(samples, size=size)
+
+    def test_broadcast_dist_samples_to(self, samples_to_broadcast_to):
+        to_shape, size, samples, broadcast_shape = samples_to_broadcast_to
+        if broadcast_shape is not None:
+            outs = broadcast_dist_samples_to(to_shape, samples, size=size)
+            assert all((o.shape == broadcast_shape for o in outs))
+        else:
+            with pytest.raises(ValueError):
+                broadcast_dist_samples_to(to_shape, samples, size=size)

From 392043072c8885e9efab5904adff9d18353237cf Mon Sep 17 00:00:00 2001
From: lucianopaz <luciano.paz.neuro@gmail.com>
Date: Fri, 26 Apr 2019 11:38:00 +0200
Subject: [PATCH 11/24] Made all tests use fixtures. Finished shape
 broadcasting tests.

---
 pymc3/tests/test_shape_handling.py | 121 ++++++++++++++++-------------
 1 file changed, 66 insertions(+), 55 deletions(-)

diff --git a/pymc3/tests/test_shape_handling.py b/pymc3/tests/test_shape_handling.py
index 7dd6e6fa0e..b35861b99f 100644
--- a/pymc3/tests/test_shape_handling.py
+++ b/pymc3/tests/test_shape_handling.py
@@ -15,18 +15,14 @@
     (tuple(), (1,), (4,), (5, 4)),
     (tuple(), (1,), (7,), (5, 4)),
     (tuple(), (1,), (1, 4), (5, 4)),
+    (tuple(), (1,), (5, 1), (5, 4)),
+    (tuple(), (1,), (3, 4), (5, 4)),
+    (tuple(), (1,), (5, 3), (5, 4)),
     (tuple(), (1,), (10, 4), (5, 4)),
     (tuple(), (1,), (10,), (5, 4)),
-    (tuple(), (1,), (1, 4), (5, 4)),
-    (tuple(), (1,), (3, 4), (5, 4)),
     (tuple(), (1,), (1, 1, 4), (5, 4)),
     (tuple(), (1,), (10, 1, 4), (5, 4)),
-    (tuple(), (1,), (5, 1), (5, 4)),
-    (tuple(), (1,), (5, 3), (5, 4)),
-    (tuple(), (1,), (10, 5, 1), (5, 4)),
-    (tuple(), (1,), (10, 5, 3), (5, 4)),
     (tuple(), (1,), (10, 5, 4), (5, 4)),
-    (tuple(), (1,), (5, 4), (5, 4)),
 ]
 test_sizes = [
     None,
@@ -37,11 +33,53 @@
     (10,),
     (1, 1),
     (10, 1),
-    (7,),
     (1, 10),
     (5,),
     (5, 4),
+    (1, 1, 1, 1),
 ]
+test_to_shapes = [None, tuple(), (10, 5, 4), (10, 1, 1, 5, 1)]
+
+
+@pytest.fixture(scope="module", params=test_sizes, ids=str)
+def fixture_sizes(request):
+    return request.param
+
+
+@pytest.fixture(scope="module", params=test_shapes, ids=str)
+def fixture_shapes(request):
+    return request.param
+
+
+@pytest.fixture(scope="module", params=[False, True], ids=str)
+def fixture_exception_handling(request):
+    return request.param
+
+
+@pytest.fixture(scope="module")
+def samples_to_broadcast(fixture_sizes, fixture_shapes):
+    samples = [np.empty(s) for s in fixture_shapes]
+    try:
+        broadcast_shape = broadcast_dist_samples_shape(
+            fixture_shapes, size=fixture_sizes
+        )
+    except ValueError:
+        broadcast_shape = None
+    return fixture_sizes, samples, broadcast_shape
+
+
+@pytest.fixture(scope="module", params=test_to_shapes, ids=str)
+def samples_to_broadcast_to(request, samples_to_broadcast):
+    to_shape = request.param
+    size, samples, broadcast_shape = samples_to_broadcast
+    if broadcast_shape is not None:
+        try:
+            broadcast_shape = broadcast_dist_samples_shape(
+                [broadcast_shape, to_tuple(to_shape)], size=size
+            )
+        except ValueError:
+            broadcast_shape = None
+    return to_shape, size, samples, broadcast_shape
 
 
 class TestShapesBroadcasting:
@@ -55,18 +93,14 @@ def test_type_check_raises(self, bad_input):
             shapes_broadcasting(bad_input, tuple(), raise_exception=True)
             shapes_broadcasting(bad_input, tuple(), raise_exception=False)
 
-    @pytest.mark.parametrize("raise_exception", [False, True], ids=str)
-    def test_type_check_success(self, raise_exception):
+    def test_type_check_success(self):
         inputs = [3, 3.0, tuple(), [3], (3,), np.array(3), np.array([3])]
-        out = shapes_broadcasting(*inputs, raise_exception=raise_exception)
+        out = shapes_broadcasting(*inputs)
         assert out == (3,)
 
-    @pytest.mark.parametrize(
-        ["shapes", "raise_exception"],
-        itertools.product(test_shapes, [False, True]),
-        ids=str,
-    )
-    def test_broadcasting(self, shapes, raise_exception):
+    def test_broadcasting(self, fixture_shapes, fixture_exception_handling):
+        shapes = fixture_shapes
+        raise_exception = fixture_exception_handling
         try:
             expected_out = np.broadcast(*[np.empty(s) for s in shapes]).shape
         except ValueError:
@@ -82,10 +116,9 @@ def test_broadcasting(self, shapes, raise_exception):
             out = shapes_broadcasting(*shapes, raise_exception=raise_exception)
             assert out == expected_out
 
-    @pytest.mark.parametrize(
-        ["size", "shapes"], itertools.product(test_sizes, test_shapes), ids=str
-    )
-    def test_broadcast_dist_samples_shape(self, size, shapes):
+    def test_broadcast_dist_samples_shape(self, fixture_sizes, fixture_shapes):
+        size = fixture_sizes
+        shapes = fixture_shapes
         size_ = to_tuple(size)
         shapes_ = [
             s if s[: min([len(size_), len(s)])] != size_ else s[len(size_) :]
@@ -107,35 +140,6 @@ def test_broadcast_dist_samples_shape(self, size, shapes):
             assert out == expected_out
 
 
-@pytest.fixture(
-    scope="module", params=itertools.product(test_sizes, test_shapes), ids=str
-)
-def samples_to_broadcast(request):
-    size, shapes = request.param
-    samples = [np.empty(s) for s in shapes]
-    try:
-        broadcast_shape = broadcast_dist_samples_shape(shapes, size=size)
-    except ValueError:
-        broadcast_shape = None
-    return size, samples, broadcast_shape
-
-
-@pytest.fixture(
-    scope="module", params=(tuple(), (1,), (10, 5, 4), (10, 1, 1, 5, 1)), ids=str
-)
-def samples_to_broadcast_to(request, samples_to_broadcast):
-    to_shape = request.param
-    size, samples, broadcast_shape = samples_to_broadcast
-    if broadcast_shape is not None:
-        try:
-            broadcast_shape = broadcast_dist_samples_shape(
-                [broadcast_shape, to_shape], size=size
-            )
-        except ValueError:
-            broadcast_shape = None
-    return to_shape, size, samples, broadcast_shape
-
-
 class TestSamplesBroadcasting:
     def test_broadcast_distribution_samples(self, samples_to_broadcast):
         size, samples, broadcast_shape = samples_to_broadcast
@@ -149,15 +153,22 @@ def test_broadcast_distribution_samples(self, samples_to_broadcast):
     def test_get_broadcastable_dist_samples(self, samples_to_broadcast):
         size, samples, broadcast_shape = samples_to_broadcast
         if broadcast_shape is not None:
+            size_ = to_tuple(size)
             outs, out_shape = get_broadcastable_dist_samples(
                 samples, size=size, return_out_shape=True
             )
             assert out_shape == broadcast_shape
-            for o in outs:
-                for oshape, bshape in itertools.zip_longest(
-                    reversed(o.shape), reversed(broadcast_shape), fillvalue=1
-                ):
-                    assert oshape == bshape or oshape == 1
+            for i, o in zip(samples, outs):
+                ishape = i.shape
+                if ishape[:min([len(size_), len(ishape)])] == size_:
+                    expected_shape = (
+                        size_ +
+                        (1,) * (len(broadcast_shape) - len(ishape)) +
+                        ishape[len(size_):]
+                    )
+                else:
+                    expected_shape = ishape
+                assert o.shape == expected_shape
             assert shapes_broadcasting(*[o.shape for o in outs]) == broadcast_shape
         else:
             with pytest.raises(ValueError):

From 95719d5df4f0e9a6ac6e21f5de042e6a3553f231 Mon Sep 17 00:00:00 2001
From: lucianopaz <luciano.paz.neuro@gmail.com>
Date: Fri, 26 Apr 2019 11:51:30 +0200
Subject: [PATCH 12/24] Fixed shape_utils docstrings.

---
 pymc3/distributions/shape_utils.py | 62 ++++++++++++++++++++++--------
 1 file changed, 46 insertions(+), 16 deletions(-)

diff --git a/pymc3/distributions/shape_utils.py b/pymc3/distributions/shape_utils.py
index 47b5da3098..6dbfe9c17f 100644
--- a/pymc3/distributions/shape_utils.py
+++ b/pymc3/distributions/shape_utils.py
@@ -38,14 +38,21 @@ def _check_shape_type(shape):
 def shapes_broadcasting(*args, raise_exception=False):
     """Return the shape resulting from broadcasting multiple shapes.
     Represents numpy's broadcasting rules.
+
     Parameters
     ----------
     *args : array-like of int
         Tuples or arrays or lists representing the shapes of arrays to be
         broadcast.
+    raise_exception: bool (optional)
+        Controls whether to raise an exception or simply return `None` if
+        the broadcasting fails.
+
     Returns
     -------
-    Resulting shape or None if broadcasting is not possible.
+    Resulting shape. If broadcasting is not possible and `raise_exception` is
+    False, then `None` is returned. If `raise_exception` is `True`, a
+    `ValueError` is raised.
     """
     x = list(_check_shape_type(args[0])) if args else ()
     for arg in args[1:]:
@@ -70,9 +77,11 @@ def shapes_broadcasting(*args, raise_exception=False):
 
 
 def broadcast_dist_samples_shape(shapes, size=None):
-    """Get the resulting broadcasted shape for samples drawn from distributions,
-    taking into account the size (i.e. the number of samples) of the draw,
-    which is prepended to the sample's shape.
+    """Apply shape broadcasting to shape tuples but assuming that the shapes
+    correspond to draws from random variables, with the `size` tuple possibly
+    prepended to it. The `size` prepend is ignored to consider if the supplied
+    `shapes` can broadcast or not. It is prepended to the resulting broadcasted
+    `shapes`, if any of the shape tuples had the `size` prepend.
 
     Parameters
     ----------
@@ -92,7 +101,7 @@ def broadcast_dist_samples_shape(shapes, size=None):
         shape1 = (size, 5)
         shape2 = (size, 4, 5)
         out = broadcast_dist_samples_shape([shape0, shape1, shape2],
-                                                   size=size)
+                                           size=size)
         assert out == (size, 4, 5)
 
     .. code-block:: python
@@ -101,8 +110,17 @@ def broadcast_dist_samples_shape(shapes, size=None):
         shape1 = (5,)
         shape2 = (4, 5)
         out = broadcast_dist_samples_shape([shape0, shape1, shape2],
-                                                   size=size)
+                                           size=size)
         assert out == (size, 4, 5)
+
+    .. code-block:: python
+        size = 100
+        shape0 = (1,)
+        shape1 = (5,)
+        shape2 = (4, 5)
+        out = broadcast_dist_samples_shape([shape0, shape1, shape2],
+                                           size=size)
+        assert out == (4, 5)
     """
     if size is None:
         broadcasted_shape = shapes_broadcasting(*shapes)
@@ -116,7 +134,10 @@ def broadcast_dist_samples_shape(shapes, size=None):
     shapes = [_check_shape_type(s) for s in shapes]
     _size = to_tuple(size)
     # samples shapes without the size prepend
-    sp_shapes = [s[len(_size) :] if _size == s[: min([len(_size), len(s)])] else s for s in shapes]
+    sp_shapes = [
+        s[len(_size) :] if _size == s[: min([len(_size), len(s)])] else s
+        for s in shapes
+    ]
     try:
         broadcast_shape = shapes_broadcasting(*sp_shapes, raise_exception=True)
     except ValueError:
@@ -152,11 +173,13 @@ def get_broadcastable_dist_samples(
     samples, size=None, must_bcast_with=None, return_out_shape=False
 ):
     """Get a view of the samples drawn from distributions which adds new axises
-    in between the size prepend and the distribution shape. These views should
-    be able to broadcast the samples from the distrubtions taking into account
-    the size (i.e. the number of samples) of the draw, which is prepended to
-    the sample's shape. Optionally, one can supply an extra `must_bcast_with`
-    to try to force samples to be able to broadcast with a given shape.
+    in between the `size` prepend and the distribution's `shape`. These views
+    should be able to broadcast the samples from the distrubtions taking into
+    account the `size` (i.e. the number of samples) of the draw, which is
+    prepended to the sample's `shape`. Optionally, one can supply an extra
+    `must_bcast_with` to try to force samples to be able to broadcast with a
+    given shape. A `ValueError` is raised if it is not possible to broadcast
+    the provided samples.
 
     Parameters
     ----------
@@ -188,7 +211,9 @@ def get_broadcastable_dist_samples(
             size=size,
             must_bcast_with=must_bcast_with,
         )
-        assert np.all((o.shape == (size, 3, 4, 5) for o in out))
+        assert out[0].shape == (size, 1, 1, 1)
+        assert out[1].shape == (size, 1, 1, 5)
+        assert out[2].shape == (size, 1, 4, 5)
         assert np.all(sample0[:, None, None, None] == out[0])
         assert np.all(sample1[:, None, None] == out[1])
         assert np.all(sample2[:, None] == out[2])
@@ -204,7 +229,9 @@ def get_broadcastable_dist_samples(
             size=size,
             must_bcast_with=must_bcast_with,
         )
-        assert np.all((o.shape == (size, 3, 4, 5) for o in out))
+        assert out[0].shape == (size, 1, 1, 1)
+        assert out[1].shape == (5,)
+        assert out[2].shape == (4, 5)
         assert np.all(sample0[:, None, None, None] == out[0])
         assert np.all(sample1 == out[1])
         assert np.all(sample2 == out[2])
@@ -216,7 +243,10 @@ def get_broadcastable_dist_samples(
     p_shapes = [p.shape for p in samples] + [_check_shape_type(must_bcast_with)]
     out_shape = broadcast_dist_samples_shape(p_shapes, size=size)
     # samples shapes without the size prepend
-    sp_shapes = [s[len(_size) :] if _size == s[: min([len(_size), len(s)])] else s for s in p_shapes]
+    sp_shapes = [
+        s[len(_size) :] if _size == s[: min([len(_size), len(s)])] else s
+        for s in p_shapes
+    ]
     broadcast_shape = shapes_broadcasting(*sp_shapes, raise_exception=True)
     broadcastable_samples = []
     for param, p_shape, sp_shape in zip(samples, p_shapes, sp_shapes):
@@ -290,7 +320,7 @@ def broadcast_dist_samples_to(to_shape, samples, size=None):
 
     Parameters
     ----------
-    to_shape: Tuple shape to which to broadcast the samples
+    to_shape: Tuple shape onto which the samples must be able to broadcast
     samples: Iterable of ndarrays holding the sampled values
     size: None, int or tuple (optional)
         size of the sample set requested.

From aebacc7f1e5bfbe886cce94330ca069930c0368e Mon Sep 17 00:00:00 2001
From: lucianopaz <luciano.paz.neuro@gmail.com>
Date: Fri, 26 Apr 2019 12:46:47 +0200
Subject: [PATCH 13/24] Added random variable sampling to test_shape testsuite.

---
 pymc3/distributions/distribution.py |  2 +-
 pymc3/tests/test_shape_handling.py  | 47 ++++++++++++++++++++++++++---
 2 files changed, 44 insertions(+), 5 deletions(-)

diff --git a/pymc3/distributions/distribution.py b/pymc3/distributions/distribution.py
index c3a32e90f2..88869012b7 100644
--- a/pymc3/distributions/distribution.py
+++ b/pymc3/distributions/distribution.py
@@ -676,7 +676,7 @@ def generate_samples(generator, *args, **kwargs):
     samples = np.asarray(samples)
 
     # reshape samples here
-    if samples.shape[0] == 1 and size == 1:
+    if samples.ndim > 0 and samples.shape[0] == 1 and size_tup == (1,):
         if (len(samples.shape) > len(dist_shape) and
             samples.shape[-len(dist_shape):] == dist_shape[-len(dist_shape):]
         ):
diff --git a/pymc3/tests/test_shape_handling.py b/pymc3/tests/test_shape_handling.py
index b35861b99f..a9673ec613 100644
--- a/pymc3/tests/test_shape_handling.py
+++ b/pymc3/tests/test_shape_handling.py
@@ -82,6 +82,33 @@ def samples_to_broadcast_to(request, samples_to_broadcast):
     return to_shape, size, samples, broadcast_shape
 
 
+@pytest.fixture(scope="module")
+def fixture_model():
+    with pm.Model() as model:
+        n = 5
+        dim = 4
+        with pm.Model():
+            cov = pm.InverseGamma("cov", alpha=1, beta=1)
+            x = pm.Normal(
+                "x", mu=np.ones((dim,)), sigma=pm.math.sqrt(cov), shape=(n, dim)
+            )
+            eps = pm.HalfNormal("eps", np.ones((n, 1)), shape=(n, dim))
+    return model, cov, x, eps
+
+
+# TODO: once #3422 is solved this fixture should be replaced by fixture_sizes
+@pytest.fixture(
+    scope="module",
+    params=[
+        s if len(to_tuple(s)) <= 1 else pytest.param(s, marks=pytest.mark.skip)
+        for s in test_sizes
+    ],
+    ids=str,
+)
+def fixture_samples(request):
+    return request.param
+
+
 class TestShapesBroadcasting:
     @pytest.mark.parametrize(
         "bad_input",
@@ -160,11 +187,11 @@ def test_get_broadcastable_dist_samples(self, samples_to_broadcast):
             assert out_shape == broadcast_shape
             for i, o in zip(samples, outs):
                 ishape = i.shape
-                if ishape[:min([len(size_), len(ishape)])] == size_:
+                if ishape[: min([len(size_), len(ishape)])] == size_:
                     expected_shape = (
-                        size_ +
-                        (1,) * (len(broadcast_shape) - len(ishape)) +
-                        ishape[len(size_):]
+                        size_
+                        + (1,) * (len(broadcast_shape) - len(ishape))
+                        + ishape[len(size_) :]
                     )
                 else:
                     expected_shape = ishape
@@ -182,3 +209,15 @@ def test_broadcast_dist_samples_to(self, samples_to_broadcast_to):
         else:
             with pytest.raises(ValueError):
                 broadcast_dist_samples_to(to_shape, samples, size=size)
+
+
+def test_sample_generate_values(fixture_model, fixture_samples):
+    model, cov, x, eps = fixture_model
+    size = to_tuple(fixture_samples)
+    if size == (1,):
+        # Single draws are interpreted as scalars for backwards compatibility
+        size = tuple()
+    with model:
+        prior = pm.sample_prior_predictive(samples=fixture_samples)
+        for rv in [cov, x, eps]:
+            assert prior[rv.name].shape == size + tuple(rv.distribution.shape)

From 183c7425d8511d4b8e6544bd583b76f24f8cc6bc Mon Sep 17 00:00:00 2001
From: lucianopaz <luciano.paz.neuro@gmail.com>
Date: Fri, 26 Apr 2019 12:50:55 +0200
Subject: [PATCH 14/24] Removed redundant broadcast_distribution_samples from
 random methods. This is now handled in generate_samples.

---
 pymc3/distributions/continuous.py | 10 ----------
 pymc3/distributions/discrete.py   |  1 -
 2 files changed, 11 deletions(-)

diff --git a/pymc3/distributions/continuous.py b/pymc3/distributions/continuous.py
index abfda57044..0e2b7a204a 100644
--- a/pymc3/distributions/continuous.py
+++ b/pymc3/distributions/continuous.py
@@ -20,7 +20,6 @@
     normal_lccdf, normal_lcdf, SplineWrapper, std_cdf, zvalue,
 )
 from .distribution import (Continuous, draw_values, generate_samples)
-from .shape_utils import broadcast_distribution_samples
 
 __all__ = ['Uniform', 'Flat', 'HalfFlat', 'Normal', 'TruncatedNormal', 'Beta',
            'Kumaraswamy', 'Exponential', 'Laplace', 'StudentT', 'Cauchy',
@@ -966,8 +965,6 @@ def random(self, point=None, size=None):
         """
         mu, lam, alpha = draw_values([self.mu, self.lam, self.alpha],
                                      point=point, size=size)
-        mu, lam, alpha = broadcast_distribution_samples([mu, lam, alpha],
-                                                        size=size)
         return generate_samples(self._random,
                                 mu, lam, alpha,
                                 dist_shape=self.shape,
@@ -1297,7 +1294,6 @@ def random(self, point=None, size=None):
         """
         a, b = draw_values([self.a, self.b],
                            point=point, size=size)
-        a, b = broadcast_distribution_samples([a, b], size=size)
         return generate_samples(self._random, a, b,
                                 dist_shape=self.shape,
                                 size=size)
@@ -1674,7 +1670,6 @@ def random(self, point=None, size=None):
         array
         """
         mu, tau = draw_values([self.mu, self.tau], point=point, size=size)
-        mu, tau = broadcast_distribution_samples([mu, tau], size=size)
         return generate_samples(self._random, mu, tau,
                                 dist_shape=self.shape,
                                 size=size)
@@ -1965,7 +1960,6 @@ def random(self, point=None, size=None):
         """
         alpha, m = draw_values([self.alpha, self.m],
                                point=point, size=size)
-        alpha, m = broadcast_distribution_samples([alpha, m], size=size)
         return generate_samples(self._random, alpha, m,
                                 dist_shape=self.shape,
                                 size=size)
@@ -2090,7 +2084,6 @@ def random(self, point=None, size=None):
         """
         alpha, beta = draw_values([self.alpha, self.beta],
                                   point=point, size=size)
-        alpha, beta = broadcast_distribution_samples([alpha, beta], size=size)
         return generate_samples(self._random, alpha, beta,
                                 dist_shape=self.shape,
                                 size=size)
@@ -2669,7 +2662,6 @@ def random(self, point=None, size=None):
         """
         alpha, beta = draw_values([self.alpha, self.beta],
                                   point=point, size=size)
-        alpha, beta = broadcast_distribution_samples([alpha, beta], size=size)
 
         def _random(a, b, size=None):
             return b * (-np.log(np.random.uniform(size=size)))**(1 / a)
@@ -2963,8 +2955,6 @@ def random(self, point=None, size=None):
         """
         mu, sigma, nu = draw_values([self.mu, self.sigma, self.nu],
                                     point=point, size=size)
-        mu, sigma, nu = broadcast_distribution_samples([mu, sigma, nu],
-                                                       size=size)
 
         def _random(mu, sigma, nu, size=None):
             return (np.random.normal(mu, sigma, size=size)
diff --git a/pymc3/distributions/discrete.py b/pymc3/distributions/discrete.py
index a2a243b3f1..b431c23458 100644
--- a/pymc3/distributions/discrete.py
+++ b/pymc3/distributions/discrete.py
@@ -351,7 +351,6 @@ def _random(self, q, beta, size=None):
 
     def random(self, point=None, size=None):
         q, beta = draw_values([self.q, self.beta], point=point, size=size)
-        q, beta = broadcast_distribution_samples([q, beta], size=size)
 
         return generate_samples(self._random, q, beta,
                                 dist_shape=self.shape,

From fac33616f09d47a6fc4a4b556f8a6dffc6712810 Mon Sep 17 00:00:00 2001
From: lucianopaz <luciano.paz.neuro@gmail.com>
Date: Fri, 26 Apr 2019 13:19:44 +0200
Subject: [PATCH 15/24] Fixed lint

---
 pymc3/tests/test_shape_handling.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/pymc3/tests/test_shape_handling.py b/pymc3/tests/test_shape_handling.py
index a9673ec613..4fedd756e9 100644
--- a/pymc3/tests/test_shape_handling.py
+++ b/pymc3/tests/test_shape_handling.py
@@ -1,5 +1,4 @@
 import pytest
-import itertools
 import numpy as np
 import pymc3 as pm
 from pymc3.distributions.shape_utils import (

From 854407c6f97d3ee6255c9825e5ada69e2624dab4 Mon Sep 17 00:00:00 2001
From: lucianopaz <luciano.paz.neuro@gmail.com>
Date: Fri, 26 Apr 2019 13:22:37 +0200
Subject: [PATCH 16/24] Moved test_shape_handling to last travis environment

---
 .travis.yml | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/.travis.yml b/.travis.yml
index c6f6d3c07a..432216fb3a 100644
--- a/.travis.yml
+++ b/.travis.yml
@@ -22,14 +22,14 @@ install:
   - conda list && pip freeze
 
 env:
-  - FLOATX='float32' TESTCMD="--durations=10 --ignore=pymc3/tests/test_examples.py --cov-append --ignore=pymc3/tests/test_distributions_random.py --ignore=pymc3/tests/test_variational_inference.py --ignore=pymc3/tests/test_shared.py --ignore=pymc3/tests/test_smc.py --ignore=pymc3/tests/test_updates.py --ignore=pymc3/tests/test_posteriors.py --ignore=pymc3/tests/test_sampling.py --ignore=pymc3/tests/test_parallel_sampling.py --ignore=pymc3/tests/test_dist_math.py --ignore=pymc3/tests/test_distribution_defaults.py --ignore=pymc3/tests/test_distributions_timeseries.py --ignore=pymc3/tests/test_random.py --ignore=pymc3/tests/test_gp.py"
+  - FLOATX='float32' TESTCMD="--durations=10 --ignore=pymc3/tests/test_examples.py --cov-append --ignore=pymc3/tests/test_distributions_random.py --ignore=pymc3/tests/test_variational_inference.py --ignore=pymc3/tests/test_shared.py --ignore=pymc3/tests/test_smc.py --ignore=pymc3/tests/test_updates.py --ignore=pymc3/tests/test_posteriors.py --ignore=pymc3/tests/test_sampling.py --ignore=pymc3/tests/test_parallel_sampling.py --ignore=pymc3/tests/test_dist_math.py --ignore=pymc3/tests/test_distribution_defaults.py --ignore=pymc3/tests/test_distributions_timeseries.py --ignore=pymc3/tests/test_random.py --ignore=pymc3/tests/test_gp.py --ignore=pymc3/tests/test_shape_handling.py"
   - FLOATX='float32' RUN_PYLINT="true" TESTCMD="--durations=10 --cov-append pymc3/tests/test_distributions_random.py pymc3/tests/test_shared.py pymc3/tests/test_smc.py pymc3/tests/test_sampling.py pymc3/tests/test_parallel_sampling.py pymc3/tests/test_dist_math.py pymc3/tests/test_distribution_defaults.py pymc3/tests/test_distributions_timeseries.py pymc3/tests/test_random.py"
   - FLOATX='float32' TESTCMD="--durations=10 --cov-append pymc3/tests/test_examples.py pymc3/tests/test_posteriors.py pymc3/tests/test_gp.py"
-  - FLOATX='float32' TESTCMD="--durations=10 --cov-append pymc3/tests/test_variational_inference.py pymc3/tests/test_updates.py"
-  - FLOATX='float64' TESTCMD="--durations=10 --cov-append --ignore=pymc3/tests/test_examples.py --ignore=pymc3/tests/test_distributions_random.py --ignore=pymc3/tests/test_variational_inference.py --ignore=pymc3/tests/test_shared.py --ignore=pymc3/tests/test_smc.py --ignore=pymc3/tests/test_updates.py --ignore=pymc3/tests/test_posteriors.py --ignore=pymc3/tests/test_sampling.py --ignore=pymc3/tests/test_parallel_sampling.py --ignore=pymc3/tests/test_dist_math.py --ignore=pymc3/tests/test_distribution_defaults.py --ignore=pymc3/tests/test_distributions_timeseries.py --ignore=pymc3/tests/test_random.py --ignore=pymc3/tests/test_gp.py"
+  - FLOATX='float32' TESTCMD="--durations=10 --cov-append pymc3/tests/test_variational_inference.py pymc3/tests/test_updates.py pymc3/tests/test_shape_handling.py"
+  - FLOATX='float64' TESTCMD="--durations=10 --cov-append --ignore=pymc3/tests/test_examples.py --ignore=pymc3/tests/test_distributions_random.py --ignore=pymc3/tests/test_variational_inference.py --ignore=pymc3/tests/test_shared.py --ignore=pymc3/tests/test_smc.py --ignore=pymc3/tests/test_updates.py --ignore=pymc3/tests/test_posteriors.py --ignore=pymc3/tests/test_sampling.py --ignore=pymc3/tests/test_parallel_sampling.py --ignore=pymc3/tests/test_dist_math.py --ignore=pymc3/tests/test_distribution_defaults.py --ignore=pymc3/tests/test_distributions_timeseries.py --ignore=pymc3/tests/test_random.py --ignore=pymc3/tests/test_gp.py --ignore=pymc3/tests/test_shape_handling.py"
   - FLOATX='float64' TESTCMD="--durations=10 --cov-append pymc3/tests/test_distributions_random.py pymc3/tests/test_shared.py pymc3/tests/test_smc.py pymc3/tests/test_sampling.py pymc3/tests/test_parallel_sampling.py pymc3/tests/test_dist_math.py pymc3/tests/test_distribution_defaults.py pymc3/tests/test_distributions_timeseries.py pymc3/tests/test_random.py"
   - FLOATX='float64' TESTCMD="--durations=10 --cov-append pymc3/tests/test_examples.py pymc3/tests/test_posteriors.py pymc3/tests/test_gp.py"
-  - FLOATX='float64' TESTCMD="--durations=10 --cov-append pymc3/tests/test_variational_inference.py pymc3/tests/test_updates.py"
+  - FLOATX='float64' TESTCMD="--durations=10 --cov-append pymc3/tests/test_variational_inference.py pymc3/tests/test_updates.py pymc3/tests/test_shape_handling.py"
 
 script:
   - . ./scripts/test.sh $TESTCMD

From d89f55c63bc3cbdc1a11828e4f684430bf9da712 Mon Sep 17 00:00:00 2001
From: lucianopaz <luciano.paz.neuro@gmail.com>
Date: Fri, 26 Apr 2019 13:31:27 +0200
Subject: [PATCH 17/24] Changed test skip to xfail

---
 pymc3/tests/test_shape_handling.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/pymc3/tests/test_shape_handling.py b/pymc3/tests/test_shape_handling.py
index 4fedd756e9..0ff0aa496d 100644
--- a/pymc3/tests/test_shape_handling.py
+++ b/pymc3/tests/test_shape_handling.py
@@ -99,7 +99,7 @@ def fixture_model():
 @pytest.fixture(
     scope="module",
     params=[
-        s if len(to_tuple(s)) <= 1 else pytest.param(s, marks=pytest.mark.skip)
+        s if len(to_tuple(s)) <= 1 else pytest.param(s, marks=pytest.mark.xfail)
         for s in test_sizes
     ],
     ids=str,

From acc6352944be2ed3820a84a1c00d1ab364b7a3d0 Mon Sep 17 00:00:00 2001
From: Luciano Paz <chulo1787@gmail.com>
Date: Fri, 26 Apr 2019 15:36:08 +0200
Subject: [PATCH 18/24] Fixed Bound random error

---
 pymc3/distributions/bound.py | 23 +++++++++++++++++++----
 1 file changed, 19 insertions(+), 4 deletions(-)

diff --git a/pymc3/distributions/bound.py b/pymc3/distributions/bound.py
index 39da708698..cffc3cc819 100644
--- a/pymc3/distributions/bound.py
+++ b/pymc3/distributions/bound.py
@@ -59,7 +59,7 @@ def _random(self, lower, upper, point=None, size=None):
                 "Drawing samples from distributions with "
                 "array-valued bounds is not supported."
             )
-        total_size = np.prod(size)
+        total_size = np.prod(size).astype(np.int)
         samples = []
         s = 0
         while s < total_size:
@@ -81,17 +81,32 @@ def random(self, point=None, size=None):
         elif self.lower is not None and self.upper is not None:
             lower, upper = draw_values([self.lower, self.upper], point=point, size=size)
             return generate_samples(
-                self._random, lower, upper, point, dist_shape=self.shape, size=size
+                self._random,
+                lower,
+                upper,
+                dist_shape=self.shape,
+                size=size,
+                not_broadcast_kwargs={'point': point},
             )
         elif self.lower is not None:
             lower = draw_values([self.lower], point=point, size=size)
             return generate_samples(
-                self._random, lower, np.inf, point, dist_shape=self.shape, size=size
+                self._random,
+                lower,
+                np.inf,
+                dist_shape=self.shape,
+                size=size,
+                not_broadcast_kwargs={'point': point},
             )
         else:
             upper = draw_values([self.upper], point=point, size=size)
             return generate_samples(
-                self._random, -np.inf, upper, point, dist_shape=self.shape, size=size
+                self._random,
+                -np.inf,
+                upper,
+                dist_shape=self.shape,
+                size=size,
+                not_broadcast_kwargs={'point': point},
             )
 
 

From 61865b73faed56319b0f682a38f09c12add7f63c Mon Sep 17 00:00:00 2001
From: lucianopaz <luciano.paz.neuro@gmail.com>
Date: Fri, 26 Apr 2019 16:59:59 +0200
Subject: [PATCH 19/24] Changes to increase code coverage

---
 pymc3/distributions/shape_utils.py | 9 +--------
 pymc3/tests/test_shape_handling.py | 1 +
 2 files changed, 2 insertions(+), 8 deletions(-)

diff --git a/pymc3/distributions/shape_utils.py b/pymc3/distributions/shape_utils.py
index 6dbfe9c17f..b799ef0d61 100644
--- a/pymc3/distributions/shape_utils.py
+++ b/pymc3/distributions/shape_utils.py
@@ -159,14 +159,7 @@ def broadcast_dist_samples_shape(shapes, size=None):
         else:
             p_shape = shape
         broadcastable_shapes.append(p_shape)
-    broadcasted_shape = shapes_broadcasting(*broadcastable_shapes)
-    if broadcasted_shape is None:
-        raise ValueError(
-            "Cannot broadcast provided shapes {} given size: {}".format(
-                ", ".join(["{}".format(s) for s in shapes]), size
-            )
-        )
-    return broadcasted_shape
+    return shapes_broadcasting(*broadcastable_shapes, raise_exception=True)
 
 
 def get_broadcastable_dist_samples(
diff --git a/pymc3/tests/test_shape_handling.py b/pymc3/tests/test_shape_handling.py
index 0ff0aa496d..5af1c8a27e 100644
--- a/pymc3/tests/test_shape_handling.py
+++ b/pymc3/tests/test_shape_handling.py
@@ -117,6 +117,7 @@ class TestShapesBroadcasting:
     def test_type_check_raises(self, bad_input):
         with pytest.raises(TypeError):
             shapes_broadcasting(bad_input, tuple(), raise_exception=True)
+        with pytest.raises(TypeError):
             shapes_broadcasting(bad_input, tuple(), raise_exception=False)
 
     def test_type_check_success(self):

From 4629bb40a3f57ad1e84432f53bbadd8cb17e8f62 Mon Sep 17 00:00:00 2001
From: lucianopaz <luciano.paz.neuro@gmail.com>
Date: Sun, 28 Apr 2019 14:56:18 +0200
Subject: [PATCH 20/24] Moved to_tuple to shape_utils

---
 pymc3/distributions/dist_math.py    | 12 +-----------
 pymc3/distributions/distribution.py |  2 +-
 pymc3/distributions/mixture.py      |  4 ++--
 pymc3/distributions/multivariate.py |  3 ++-
 pymc3/distributions/shape_utils.py  | 13 ++++++++++++-
 pymc3/tests/test_mixture.py         |  2 +-
 pymc3/tests/test_shape_handling.py  |  8 ++++----
 7 files changed, 23 insertions(+), 21 deletions(-)

diff --git a/pymc3/distributions/dist_math.py b/pymc3/distributions/dist_math.py
index 044f5ad86f..d65333025f 100644
--- a/pymc3/distributions/dist_math.py
+++ b/pymc3/distributions/dist_math.py
@@ -11,6 +11,7 @@
 from theano.tensor.slinalg import Cholesky
 from theano.scan_module import until
 from theano import scan
+from .shape_utils import to_tuple
 
 from .special import gammaln
 from pymc3.theanof import floatX
@@ -20,17 +21,6 @@
 c = - .5 * np.log(2. * np.pi)
 
 
-def to_tuple(shape):
-    """Convert ints, arrays, and Nones to tuples"""
-    if shape is None:
-        return tuple()
-    temp = np.atleast_1d(shape)
-    if temp.size == 0:
-        return tuple()
-    else:
-        return tuple(temp)
-
-
 def bound(logp, *conditions, **kwargs):
     """
     Bounds a log probability density with several conditions.
diff --git a/pymc3/distributions/distribution.py b/pymc3/distributions/distribution.py
index 88869012b7..11bdab07d9 100644
--- a/pymc3/distributions/distribution.py
+++ b/pymc3/distributions/distribution.py
@@ -10,8 +10,8 @@
     ObservedRV, MultiObservedRV, Context, InitContextMeta
 )
 from ..vartypes import string_types
-from .dist_math import to_tuple
 from .shape_utils import (
+    to_tuple,
     get_broadcastable_dist_samples,
     broadcast_dist_samples_shape,
 )
diff --git a/pymc3/distributions/mixture.py b/pymc3/distributions/mixture.py
index 4bae5aeb82..b57e0386ae 100644
--- a/pymc3/distributions/mixture.py
+++ b/pymc3/distributions/mixture.py
@@ -5,11 +5,11 @@
 
 from pymc3.util import get_variable_name
 from ..math import logsumexp
-from .dist_math import bound, random_choice, to_tuple
+from .dist_math import bound, random_choice
 from .distribution import (Discrete, Distribution, draw_values,
                            generate_samples, _DrawValuesContext,
                            _DrawValuesContextBlocker)
-from .shape_utils import broadcast_distribution_samples
+from .shape_utils import to_tuple, broadcast_distribution_samples
 from .continuous import get_tau_sigma, Normal
 from ..theanof import _conversion_map
 
diff --git a/pymc3/distributions/multivariate.py b/pymc3/distributions/multivariate.py
index cefea4a7e1..2fc77446c6 100755
--- a/pymc3/distributions/multivariate.py
+++ b/pymc3/distributions/multivariate.py
@@ -21,7 +21,8 @@
 from ..model import Deterministic
 from .continuous import ChiSquared, Normal
 from .special import gammaln, multigammaln
-from .dist_math import bound, logpow, factln, to_tuple
+from .dist_math import bound, logpow, factln
+from .shape_utils import to_tuple
 from ..math import kron_dot, kron_diag, kron_solve_lower, kronecker
 
 
diff --git a/pymc3/distributions/shape_utils.py b/pymc3/distributions/shape_utils.py
index b799ef0d61..a98db90c9f 100644
--- a/pymc3/distributions/shape_utils.py
+++ b/pymc3/distributions/shape_utils.py
@@ -5,10 +5,10 @@
 samples from probability distributions for stochastic nodes in PyMC.
 """
 import numpy as np
-from .dist_math import to_tuple
 
 
 __all__ = [
+    "to_tuple",
     "shapes_broadcasting",
     "broadcast_dist_samples_shape",
     "get_broadcastable_dist_samples",
@@ -17,6 +17,17 @@
 ]
 
 
+def to_tuple(shape):
+    """Convert ints, arrays, and Nones to tuples"""
+    if shape is None:
+        return tuple()
+    temp = np.atleast_1d(shape)
+    if temp.size == 0:
+        return tuple()
+    else:
+        return tuple(temp)
+
+
 def _check_shape_type(shape):
     out = []
     try:
diff --git a/pymc3/tests/test_mixture.py b/pymc3/tests/test_mixture.py
index 45a1f6127a..ffe799f012 100644
--- a/pymc3/tests/test_mixture.py
+++ b/pymc3/tests/test_mixture.py
@@ -11,7 +11,7 @@
 from pymc3.theanof import floatX
 import theano
 from theano import tensor as tt
-from pymc3.distributions.dist_math import to_tuple
+from pymc3.distributions.shape_utils import to_tuple
 
 # Generate data
 def generate_normal_mixture_data(w, mu, sd, size=1000):
diff --git a/pymc3/tests/test_shape_handling.py b/pymc3/tests/test_shape_handling.py
index 5af1c8a27e..1d7957f69f 100644
--- a/pymc3/tests/test_shape_handling.py
+++ b/pymc3/tests/test_shape_handling.py
@@ -2,13 +2,13 @@
 import numpy as np
 import pymc3 as pm
 from pymc3.distributions.shape_utils import (
+    to_tuple,
     shapes_broadcasting,
     broadcast_dist_samples_shape,
     get_broadcastable_dist_samples,
     broadcast_distribution_samples,
     broadcast_dist_samples_to,
 )
-from pymc3.distributions.dist_math import to_tuple
 
 test_shapes = [
     (tuple(), (1,), (4,), (5, 4)),
@@ -211,13 +211,13 @@ def test_broadcast_dist_samples_to(self, samples_to_broadcast_to):
                 broadcast_dist_samples_to(to_shape, samples, size=size)
 
 
-def test_sample_generate_values(fixture_model, fixture_samples):
+def test_sample_generate_values(fixture_model, fixture_sizes):
     model, cov, x, eps = fixture_model
-    size = to_tuple(fixture_samples)
+    size = to_tuple(fixture_sizes)
     if size == (1,):
         # Single draws are interpreted as scalars for backwards compatibility
         size = tuple()
     with model:
-        prior = pm.sample_prior_predictive(samples=fixture_samples)
+        prior = pm.sample_prior_predictive(samples=fixture_sizes)
         for rv in [cov, x, eps]:
             assert prior[rv.name].shape == size + tuple(rv.distribution.shape)

From aec076e11d3c452b208e9a32dc6a87c311828c48 Mon Sep 17 00:00:00 2001
From: lucianopaz <luciano.paz.neuro@gmail.com>
Date: Sun, 28 Apr 2019 17:51:42 +0200
Subject: [PATCH 21/24] Used numpy vectorize to fix 3422

---
 RELEASE-NOTES.md                    |  1 +
 pymc3/distributions/distribution.py | 74 ++++++++++++++++++-----------
 pymc3/tests/test_shape_handling.py  |  9 ++--
 3 files changed, 53 insertions(+), 31 deletions(-)

diff --git a/RELEASE-NOTES.md b/RELEASE-NOTES.md
index 83e3540197..b824ac83ca 100644
--- a/RELEASE-NOTES.md
+++ b/RELEASE-NOTES.md
@@ -9,6 +9,7 @@
 - `Mixture` now supports mixtures of multidimensional probability distributions, not just lists of 1D distributions.
 - `GLM.from_formula` and `LinearComponent.from_formula` can extract variables from the calling scope. Customizable via the new `eval_env` argument. Fixing #3382.
 - Added the `distributions.shape_utils` module with functions used to help broadcast samples drawn from distributions using the `size` keyword argument.
+- Used `numpy.vectorize` in `distributions.distribution._compile_theano_function`. This enables `sample_prior_predictive` and `sample_posterior_predictive` to ask for tuples of samples instead of just integers. This fixes issue #3422.
 
 ### Maintenance
 - All occurances of `sd` as a parameter name have been renamed to `sigma`. `sd` will continue to function for backwards compatibility.
diff --git a/pymc3/distributions/distribution.py b/pymc3/distributions/distribution.py
index 11bdab07d9..8e65cbb265 100644
--- a/pymc3/distributions/distribution.py
+++ b/pymc3/distributions/distribution.py
@@ -451,10 +451,38 @@ def _compile_theano_function(param, vars, givens=None):
     A compiled theano function that takes the values of `vars` as input
         positional args
     """
-    return function(vars, param, givens=givens,
-                    rebuild_strict=True,
-                    on_unused_input='ignore',
-                    allow_input_downcast=True)
+    f = function(
+        vars,
+        param,
+        givens=givens,
+        rebuild_strict=True,
+        on_unused_input="ignore",
+        allow_input_downcast=True,
+    )
+    return vectorize_theano_function(f, inputs=vars, output=param)
+
+
+def vectorize_theano_function(f, inputs, output):
+    inputs_signatures = ",".join(
+        [
+            get_vectorize_signature(var, var_name="i_{}".format(input_ind))
+            for input_ind, var in enumerate(inputs)
+        ]
+    )
+    output_signature = get_vectorize_signature(output, var_name="o")
+    signature = inputs_signatures + "->" + output_signature
+
+    return np.vectorize(f, signature=signature)
+
+
+def get_vectorize_signature(var, var_name="i"):
+    if var.ndim == 0:
+        return "()"
+    else:
+        sig = ",".join(
+            ["{}_{}".format(var_name, axis_ind) for axis_ind in range(var.ndim)]
+        )
+        return "({})".format(sig)
 
 
 def _draw_value(param, point=None, givens=None, size=None):
@@ -541,19 +569,7 @@ def _draw_value(param, point=None, givens=None, size=None):
                 input_vars = []
                 input_vals = []
             func = _compile_theano_function(param, input_vars)
-            if size is not None:
-                size = np.atleast_1d(size)
-            dshaped_variables = all((hasattr(var, 'dshape')
-                                     for var in input_vars))
-            if (values and dshaped_variables and
-                not all(var.dshape == getattr(val, 'shape', tuple())
-                        for var, val in zip(input_vars, input_vals))):
-                output = np.array([func(*v) for v in zip(*input_vals)])
-            elif (size is not None and any((val.ndim > var.ndim)
-                  for var, val in zip(input_vars, input_vals))):
-                output = np.array([func(*v) for v in zip(*input_vals)])
-            else:
-                output = func(*input_vals)
+            output = func(*input_vals)
             return output
     raise ValueError('Unexpected type in draw_value: %s' % type(param))
 
@@ -615,6 +631,16 @@ def generate_samples(generator, *args, **kwargs):
     # Convert size and dist_shape to tuples
     size_tup = to_tuple(size)
     dist_shape = to_tuple(dist_shape)
+    if dist_shape[:len(size_tup)] == size_tup:
+        # dist_shape is prepended with size_tup. This is not a consequence
+        # of the parameters being drawn size_tup times! By chance, the
+        # distribution's shape has its first elements equal to size_tup.
+        # This means that we must prepend the size_tup to dist_shape, and
+        # check if that broadcasts well with the parameters
+        _dist_shape = size_tup + dist_shape
+    else:
+        _dist_shape = dist_shape
+
 
     if broadcast_shape is None:
         # If broadcast_shape is not explicitly provided, it is inferred as the
@@ -622,7 +648,7 @@ def generate_samples(generator, *args, **kwargs):
         # account the potential size prefix
         inputs = args + tuple(kwargs.values())
         broadcast_shape = broadcast_dist_samples_shape(
-            [np.asarray(i).shape for i in inputs] + [dist_shape],
+            [np.asarray(i).shape for i in inputs] + [_dist_shape],
             size=size_tup
         )
         # We do this instead of broadcast_distribution_samples to avoid
@@ -642,15 +668,6 @@ def generate_samples(generator, *args, **kwargs):
     # We ensure that broadcast_shape is a tuple
     broadcast_shape = to_tuple(broadcast_shape)
 
-    if dist_shape[:len(size_tup)] == size_tup:
-        # dist_shape is prepended with size_tup. This is not a consequence
-        # of the parameters being drawn size_tup times! By chance, the
-        # distribution's shape has its first elements equal to size_tup.
-        # This means that we must prepend the size_tup to dist_shape, and
-        # check if that broadcasts well with the parameters
-        _dist_shape = size_tup + dist_shape
-    else:
-        _dist_shape = dist_shape
     try:
         dist_bcast_shape = broadcast_dist_samples_shape(
             [_dist_shape, broadcast_shape],
@@ -682,6 +699,7 @@ def generate_samples(generator, *args, **kwargs):
         ):
             samples = samples.reshape(samples.shape[1:])
 
-    if one_d and samples.ndim > 0 and samples.shape[-1] == 1:
+    if (one_d and samples.ndim > 0 and samples.shape[-1] == 1
+        and (size_tup == tuple() or size_tup == (1,))):
         samples = samples.reshape(samples.shape[:-1])
     return np.asarray(samples)
\ No newline at end of file
diff --git a/pymc3/tests/test_shape_handling.py b/pymc3/tests/test_shape_handling.py
index 1d7957f69f..9a65bb6465 100644
--- a/pymc3/tests/test_shape_handling.py
+++ b/pymc3/tests/test_shape_handling.py
@@ -1,5 +1,6 @@
 import pytest
 import numpy as np
+from theano import tensor as tt
 import pymc3 as pm
 from pymc3.distributions.shape_utils import (
     to_tuple,
@@ -92,7 +93,9 @@ def fixture_model():
                 "x", mu=np.ones((dim,)), sigma=pm.math.sqrt(cov), shape=(n, dim)
             )
             eps = pm.HalfNormal("eps", np.ones((n, 1)), shape=(n, dim))
-    return model, cov, x, eps
+            mu = pm.Deterministic("mu", tt.sum(x + eps, axis=-1))
+            y = pm.Normal("y", mu=mu, sigma=1, shape=(n,))
+    return model, [cov, x, eps, y]
 
 
 # TODO: once #3422 is solved this fixture should be replaced by fixture_sizes
@@ -212,12 +215,12 @@ def test_broadcast_dist_samples_to(self, samples_to_broadcast_to):
 
 
 def test_sample_generate_values(fixture_model, fixture_sizes):
-    model, cov, x, eps = fixture_model
+    model, RVs = fixture_model
     size = to_tuple(fixture_sizes)
     if size == (1,):
         # Single draws are interpreted as scalars for backwards compatibility
         size = tuple()
     with model:
         prior = pm.sample_prior_predictive(samples=fixture_sizes)
-        for rv in [cov, x, eps]:
+        for rv in RVs:
             assert prior[rv.name].shape == size + tuple(rv.distribution.shape)

From 95e932a3e901991b3c254a63c1a32104b7f21dd1 Mon Sep 17 00:00:00 2001
From: lucianopaz <luciano.paz.neuro@gmail.com>
Date: Mon, 29 Apr 2019 13:13:47 +0200
Subject: [PATCH 22/24] Fixed signature for empty input and one_d reshaping.

---
 pymc3/distributions/distribution.py | 43 ++++++++++++++++++++++++++---
 1 file changed, 39 insertions(+), 4 deletions(-)

diff --git a/pymc3/distributions/distribution.py b/pymc3/distributions/distribution.py
index 8e65cbb265..ec249a713f 100644
--- a/pymc3/distributions/distribution.py
+++ b/pymc3/distributions/distribution.py
@@ -463,16 +463,48 @@ def _compile_theano_function(param, vars, givens=None):
 
 
 def vectorize_theano_function(f, inputs, output):
+    """Takes a compiled theano function and wraps it with a vectorized version.
+    Theano compiled functions expect inputs and outputs of a fixed number of
+    dimensions. In our context, these usually come from deterministics which
+    are compiled against a given RV, with its core shape. If we draw i.i.d.
+    samples from said RV, we would not be able to compute the deterministic
+    over the i.i.d sampled dimensions (i.e. those that are not the core
+    dimensions of the RV). To deal with this problem, we wrap the theano
+    compiled function with numpy.vectorize, providing the correct signature
+    for the core dimensions. The extra dimensions, will be interpreted as
+    i.i.d. sampled axis and will be broadcast following the usual rules.
+
+    Parameters
+    ----------
+    f : theano compiled function
+    inputs : list of theano variables used as inputs for the function
+    givens : theano variable which is the output of the function
+
+    Notes
+    -----
+    If inputs is an empty list (theano function with no inputs needed), then
+    the same `f` is returned.
+    Only functions that return a single theano variable's value can be
+    vectorized.
+
+    Returns
+    -------
+    A function which wraps `f` with numpy.vectorize with the apropriate call
+    signature.
+    """
     inputs_signatures = ",".join(
         [
             get_vectorize_signature(var, var_name="i_{}".format(input_ind))
             for input_ind, var in enumerate(inputs)
         ]
     )
-    output_signature = get_vectorize_signature(output, var_name="o")
-    signature = inputs_signatures + "->" + output_signature
+    if len(inputs_signatures) > 0:
+        output_signature = get_vectorize_signature(output, var_name="o")
+        signature = inputs_signatures + "->" + output_signature
 
-    return np.vectorize(f, signature=signature)
+        return np.vectorize(f, signature=signature)
+    else:
+        return f
 
 
 def get_vectorize_signature(var, var_name="i"):
@@ -700,6 +732,9 @@ def generate_samples(generator, *args, **kwargs):
             samples = samples.reshape(samples.shape[1:])
 
     if (one_d and samples.ndim > 0 and samples.shape[-1] == 1
-        and (size_tup == tuple() or size_tup == (1,))):
+        and (samples.shape != size_tup or
+             size_tup == tuple() or
+             size_tup == (1,))
+    ):
         samples = samples.reshape(samples.shape[:-1])
     return np.asarray(samples)
\ No newline at end of file

From fa83c31dfaf024402ebcced25dc12b26b1d74c4a Mon Sep 17 00:00:00 2001
From: lucianopaz <luciano.paz.neuro@gmail.com>
Date: Tue, 30 Apr 2019 15:18:34 +0200
Subject: [PATCH 23/24] Removed unused fixture and added more details to
 RELEASE-NOTES

---
 RELEASE-NOTES.md                   |  2 ++
 pymc3/tests/test_shape_handling.py | 13 -------------
 2 files changed, 2 insertions(+), 13 deletions(-)

diff --git a/RELEASE-NOTES.md b/RELEASE-NOTES.md
index b824ac83ca..89154af9e9 100644
--- a/RELEASE-NOTES.md
+++ b/RELEASE-NOTES.md
@@ -35,6 +35,8 @@
 - Resolved issue #3399. Converted all calls to `pm.distributions.bound._ContinuousBounded` and `pm.distributions.bound._DiscreteBounded` to use only and all positional arguments.
 - Restructured `distributions.distribution.generate_samples` to use the `shape_utils` module. This solves issues #3421 and #3147 by using the `size` aware broadcating functions in `shape_utils`.
 - Fixed the `Multinomial.random` and `Multinomial.random_` methods to make them compatible with the new `generate_samples` function. In the process, a bug of the `Multinomial.random_` shape handling was discovered and fixed.
+- Fixed a defect found in `Bound.random` where the `point` dictionary was passed to `generate_samples` as an `arg` instead of in `not_broadcast_kwargs`.
+- Fixed a defect found in `Bound.random_` where `total_size` could end up as a `float64` instead of being an integer if given `size=tuple()`.
 
 ### Deprecations
 
diff --git a/pymc3/tests/test_shape_handling.py b/pymc3/tests/test_shape_handling.py
index 9a65bb6465..3a1cd61375 100644
--- a/pymc3/tests/test_shape_handling.py
+++ b/pymc3/tests/test_shape_handling.py
@@ -98,19 +98,6 @@ def fixture_model():
     return model, [cov, x, eps, y]
 
 
-# TODO: once #3422 is solved this fixture should be replaced by fixture_sizes
-@pytest.fixture(
-    scope="module",
-    params=[
-        s if len(to_tuple(s)) <= 1 else pytest.param(s, marks=pytest.mark.xfail)
-        for s in test_sizes
-    ],
-    ids=str,
-)
-def fixture_samples(request):
-    return request.param
-
-
 class TestShapesBroadcasting:
     @pytest.mark.parametrize(
         "bad_input",

From fb267c5b64b1a8655bc65af63a1438b8da5d862d Mon Sep 17 00:00:00 2001
From: lucianopaz <luciano.paz.neuro@gmail.com>
Date: Thu, 2 May 2019 06:24:08 +0200
Subject: [PATCH 24/24] Addressed colcarrol comments

---
 pymc3/distributions/distribution.py | 6 +++---
 pymc3/distributions/shape_utils.py  | 1 -
 2 files changed, 3 insertions(+), 4 deletions(-)

diff --git a/pymc3/distributions/distribution.py b/pymc3/distributions/distribution.py
index ec249a713f..0164ff691d 100644
--- a/pymc3/distributions/distribution.py
+++ b/pymc3/distributions/distribution.py
@@ -709,14 +709,14 @@ def generate_samples(generator, *args, **kwargs):
         raise TypeError('''Attempted to generate values with incompatible shapes:
             size: {size}
             size_tup: {size_tup}
-            broadcast_shape[:len(size_tup)] == size_tup: {test}
+            broadcast_shape[:len(size_tup)] == size_tup: {size_prepended}
             dist_shape: {dist_shape}
             broadcast_shape: {broadcast_shape}
         '''.format(size=size,
                    size_tup=size_tup,
                    dist_shape=dist_shape,
                    broadcast_shape=broadcast_shape,
-                   test=broadcast_shape[:len(size_tup)] == size_tup)
+                   size_prepended=broadcast_shape[:len(size_tup)] == size_tup)
         )
     if dist_bcast_shape[:len(size_tup)] == size_tup:
         samples = generator(size=dist_bcast_shape, *args, **kwargs)
@@ -737,4 +737,4 @@ def generate_samples(generator, *args, **kwargs):
              size_tup == (1,))
     ):
         samples = samples.reshape(samples.shape[:-1])
-    return np.asarray(samples)
\ No newline at end of file
+    return np.asarray(samples)
diff --git a/pymc3/distributions/shape_utils.py b/pymc3/distributions/shape_utils.py
index a98db90c9f..02736a7524 100644
--- a/pymc3/distributions/shape_utils.py
+++ b/pymc3/distributions/shape_utils.py
@@ -1,4 +1,3 @@
-#!/usr/bin/env python3
 # -*- coding: utf-8 -*-
 """
 A collection of common numpy array shape operations needed for broadcasting