Make change_rv_size more robust

Co-authored-by: Ricardo <[email protected]>
Co-authored-by: Thomas Wiecki <[email protected]>

Author: 3 people

pymc3/aesaraf.py

@@ -46,10 +46,12 @@
 from aesara.sandbox.rng_mrg import MRG_RandomStream as RandomStream
 from aesara.tensor.elemwise import Elemwise
 from aesara.tensor.random.op import RandomVariable
+from aesara.tensor.shape import SpecifyShape
 from aesara.tensor.sharedvar import SharedVariable
 from aesara.tensor.subtensor import AdvancedIncSubtensor, AdvancedIncSubtensor1
 from aesara.tensor.var import TensorVariable
 
+from pymc3.exceptions import ShapeError
 from pymc3.vartypes import continuous_types, int_types, isgenerator, typefilter
 
 PotentialShapeType = Union[
@@ -153,6 +155,16 @@ def change_rv_size(
         Expand the existing size by `new_size`.
 
     """
+    # Check the dimensionality of the `new_size` kwarg
+    new_size_ndim = np.ndim(new_size)
+    if new_size_ndim > 1:
+        raise ShapeError("The `new_size` must be ≤1-dimensional.", actual=new_size_ndim)
+    elif new_size_ndim == 0:
+        new_size = (new_size,)
+
+    # Extract the RV node that is to be resized, together with its inputs, name and tag
+    if isinstance(rv_var.owner.op, SpecifyShape):
+        rv_var = rv_var.owner.inputs[0]
     rv_node = rv_var.owner
     rng, size, dtype, *dist_params = rv_node.inputs
     name = rv_var.name
@@ -161,10 +173,10 @@ def change_rv_size(
     if expand:
         if rv_node.op.ndim_supp == 0 and at.get_vector_length(size) == 0:
             size = rv_node.op._infer_shape(size, dist_params)
-        new_size = tuple(at.atleast_1d(new_size)) + tuple(size)
+        new_size = tuple(new_size) + tuple(size)
 
-    # Make sure the new size is a tensor. This helps to not unnecessarily pick
-    # up a `Cast` in some cases
+    # Make sure the new size is a tensor. This dtype-aware conversion helps
+    # to not unnecessarily pick up a `Cast` in some cases (see #4652).
     new_size = at.as_tensor(new_size, ndim=1, dtype="int64")
 
     new_rv_node = rv_node.op.make_node(rng, new_size, dtype, *dist_params)

pymc3/tests/test_aesaraf.py

@@ -41,6 +41,7 @@
     take_along_axis,
     walk_model,
 )
+from pymc3.exceptions import ShapeError
 from pymc3.vartypes import int_types
 
 FLOATX = str(aesara.config.floatX)
@@ -53,6 +54,11 @@ def test_change_rv_size():
     assert rv.ndim == 1
     assert tuple(rv.shape.eval()) == (2,)
 
+    with pytest.raises(ShapeError, match="must be ≤1-dimensional"):
+        change_rv_size(rv, new_size=[[2, 3]])
+    with pytest.raises(ShapeError, match="must be ≤1-dimensional"):
+        change_rv_size(rv, new_size=at.as_tensor_variable([[2, 3], [4, 5]]))
+
     rv_new = change_rv_size(rv, new_size=(3,), expand=True)
     assert rv_new.ndim == 2
     assert tuple(rv_new.shape.eval()) == (3, 2)

pymc3/tests/test_distributions_random.py

@@ -187,11 +187,10 @@ def get_random_variable(self, shape, with_vector_params=False, name=None):
 
         @staticmethod
         def sample_random_variable(random_variable, size):
-            """Draws samples from a RandomVariable using its .random() method."""
-            if size is None:
-                return random_variable.eval()
-            else:
-                return change_rv_size(random_variable, size, expand=True).eval()
+            """ Draws samples from a RandomVariable. """
+            if size:
+                random_variable = change_rv_size(random_variable, size, expand=True)
+            return random_variable.eval()
 
         @pytest.mark.parametrize("size", [None, (), 1, (1,), 5, (4, 5)], ids=str)
         @pytest.mark.parametrize("shape", [None, ()], ids=str)