Check ndim of new_size kwarg in change_rv_size

And don't use np.atleast_1d on things that can be tensors.

Author: michaelosthege

pymc3/aesaraf.py

@@ -50,6 +50,7 @@
 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[
@@ -147,6 +148,10 @@ def change_rv_size(
         Expand the existing size by `new_size`.
 
     """
+    new_size_ndim = new_size.ndim if isinstance(new_size, Variable) else np.ndim(new_size)
+    if new_size_ndim > 1:
+        raise ShapeError("The `new_size` must be ≤1-dimensional.", actual=new_size_ndim)
+    new_size = at.as_tensor_variable(new_size, ndim=1)
     if isinstance(rv_var.owner.op, SpecifyShape):
         rv_var = rv_var.owner.inputs[0]
     rv_node = rv_var.owner
@@ -157,7 +162,7 @@ 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(np.atleast_1d(new_size)) + tuple(size)
+        new_size = tuple(new_size) + tuple(size)
 
     new_rv_node = rv_node.op.make_node(rng, new_size, dtype, *dist_params)
     rv_var = new_rv_node.outputs[-1]

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 rv.eval().shape == (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 rv_new.eval().shape == (3, 2)