Use aesara.tensor.atleast_1d in pymc3.aesaraf.change_rv_size

Author: brandonwillard

pymc3/aesaraf.py

@@ -160,7 +160,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(at.atleast_1d(new_size)) + tuple(size)
 
     # Make sure the new size is a tensor. This helps to not unnecessarily pick
     # up a `Cast` in some cases

pymc3/tests/test_aesaraf.py

@@ -51,11 +51,11 @@ def test_change_rv_size():
     loc = at.as_tensor_variable([1, 2])
     rv = normal(loc=loc)
     assert rv.ndim == 1
-    assert rv.eval().shape == (2,)
+    assert tuple(rv.shape.eval()) == (2,)
 
     rv_new = change_rv_size(rv, new_size=(3,), expand=True)
     assert rv_new.ndim == 2
-    assert rv_new.eval().shape == (3, 2)
+    assert tuple(rv_new.shape.eval()) == (3, 2)
 
     # Make sure that the shape used to determine the expanded size doesn't
     # depend on the old `RandomVariable`.
@@ -65,7 +65,7 @@ def test_change_rv_size():
 
     rv_newer = change_rv_size(rv_new, new_size=(4,), expand=True)
     assert rv_newer.ndim == 3
-    assert rv_newer.eval().shape == (4, 3, 2)
+    assert tuple(rv_newer.shape.eval()) == (4, 3, 2)
 
     # Make sure we avoid introducing a `Cast` by converting the new size before
     # constructing the new `RandomVariable`
@@ -74,7 +74,19 @@ def test_change_rv_size():
     rv_newer = change_rv_size(rv, new_size=new_size, expand=False)
     assert rv_newer.ndim == 2
     assert isinstance(rv_newer.owner.inputs[1], Constant)
-    assert rv_newer.eval().shape == (4, 3)
+    assert tuple(rv_newer.shape.eval()) == (4, 3)
+
+    rv = normal(0, 1)
+    new_size = at.as_tensor(np.array([4, 3], dtype="int32"))
+    rv_newer = change_rv_size(rv, new_size=new_size, expand=True)
+    assert rv_newer.ndim == 2
+    assert tuple(rv_newer.shape.eval()) == (4, 3)
+
+    rv = normal(0, 1)
+    new_size = at.as_tensor(2, dtype="int32")
+    rv_newer = change_rv_size(rv, new_size=new_size, expand=True)
+    assert rv_newer.ndim == 1
+    assert tuple(rv_newer.shape.eval()) == (2,)
 
 
 class TestBroadcasting: