👌 use broadcast_dist_samples_shape

Author: symeneses

docs/source/api/shape_utils.rst

@@ -14,5 +14,6 @@ This module introduces functions that are made aware of the requested `size_tupl
    :toctree: generated/
 
    to_tuple
+   broadcast_dist_samples_shape
    rv_size_is_none
    change_dist_size

pymc/distributions/multivariate.py

@@ -59,6 +59,7 @@
 )
 from pymc.distributions.shape_utils import (
     _change_dist_size,
+    broadcast_dist_samples_shape,
     change_dist_size,
     get_support_shape,
     rv_size_is_none,
@@ -1651,8 +1652,8 @@ def rng_fn(cls, rng, mu, rowchol, colchol, size=None):
 
         # Broadcasting all parameters
         shapes = [mu.shape, output_shape]
-        sp_shapes = [s[len(size) :] if size == s[: min([len(size), len(s)])] else s for s in shapes]
-        mu = np.broadcast_to(mu, shape=np.broadcast_shapes(*sp_shapes))
+        broadcastable_shape = broadcast_dist_samples_shape(shapes, size=size)
+        mu = np.broadcast_to(mu, shape=broadcastable_shape)
         rowchol = np.broadcast_to(rowchol, shape=size + rowchol.shape[-2:])
 
         colchol = np.broadcast_to(colchol, shape=size + colchol.shape[-2:])

pymc/distributions/shape_utils.py

@@ -38,6 +38,7 @@
 from pymc.pytensorf import convert_observed_data
 
 __all__ = [
+    "broadcast_dist_samples_shape",
     "to_tuple",
     "rv_size_is_none",
     "change_dist_size",
@@ -86,45 +87,85 @@ def _check_shape_type(shape):
     return tuple(out)
 
 
-def shapes_broadcasting(*args, raise_exception=False):
-    """Return the shape resulting from broadcasting multiple shapes.
-    Represents numpy's broadcasting rules.
+def broadcast_dist_samples_shape(shapes, size=None):
+    """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
     ----------
-    *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.
+    shapes: Iterable of tuples holding the distribution samples shapes
+    size: None, int or tuple (optional)
+        size of the sample set requested.
 
     Returns
     -------
-    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.
+    tuple of the resulting shape
+
+    Examples
+    --------
+    .. code-block:: python
+        size = 100
+        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
+        shape0 = (size,)
+        shape1 = (5,)
+        shape2 = (4, 5)
+        out = broadcast_dist_samples_shape([shape0, shape1, shape2],
+                                           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)
     """
-    x = list(_check_shape_type(args[0])) if args else ()
-    for arg in args[1:]:
-        y = list(_check_shape_type(arg))
-        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)
-            ]
-        if not all(x):
-            if raise_exception:
-                raise ValueError(
-                    "Supplied shapes {} do not broadcast together".format(
-                        ", ".join([f"{a}" for a in args])
-                    )
+    if size is None:
+        broadcasted_shape = np.broadcast_shapes(*shapes)
+        if broadcasted_shape is None:
+            raise ValueError(
+                "Cannot broadcast provided shapes {} given size: {}".format(
+                    ", ".join([f"{s}" for s in shapes]), size
                 )
-            else:
-                return None
-    return tuple(x)
+            )
+        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[: min([len(_size), len(s)])] else s for s in shapes]
+    try:
+        broadcast_shape = np.broadcast_shapes(*sp_shapes)
+    except ValueError:
+        raise ValueError(
+            "Cannot broadcast provided shapes {} given size: {}".format(
+                ", ".join([f"{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)
+    return np.broadcast_shapes(*broadcastable_shapes)
 
 
 # User-provided can be lazily specified as scalars

tests/distributions/test_shape_utils.py

@@ -29,6 +29,7 @@
 
 from pymc import ShapeError
 from pymc.distributions.shape_utils import (
+    broadcast_dist_samples_shape,
     change_dist_size,
     convert_dims,
     convert_shape,
@@ -85,6 +86,43 @@ def fixture_exception_handling(request):
     return request.param
 
 
+class TestShapesBroadcasting:
+    def test_broadcasting(self, fixture_shapes):
+        shapes = fixture_shapes
+        try:
+            expected_out = np.broadcast(*(np.empty(s) for s in shapes)).shape
+        except ValueError:
+            expected_out = None
+        if expected_out is None:
+            with pytest.raises(ValueError):
+                np.broadcast_shapes(*shapes)
+        else:
+            out = np.broadcast_shapes(*shapes)
+            assert out == expected_out
+
+    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_) :] 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[: min([len(size_), len(s)])] == 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
+
+
 class TestSizeShapeDimsObserved:
     @pytest.mark.parametrize("param_shape", [(), (2,)])
     @pytest.mark.parametrize("batch_shape", [(), (3,)])