Broadcast dist in Censored with lower and upper when size is not specified

Author: ricardoV94

pymc/distributions/censored.py

@@ -18,6 +18,7 @@
 from aesara.tensor import TensorVariable
 from aesara.tensor.random.op import RandomVariable
 
+from pymc.aesaraf import change_rv_size
 from pymc.distributions.distribution import SymbolicDistribution, _moment
 from pymc.util import check_dist_not_registered
 
@@ -74,10 +75,13 @@ def dist(cls, dist, lower, upper, **kwargs):
 
     @classmethod
     def rv_op(cls, dist, lower=None, upper=None, size=None, rngs=None):
-        if lower is None:
-            lower = at.constant(-np.inf)
-        if upper is None:
-            upper = at.constant(np.inf)
+
+        lower = at.constant(-np.inf) if lower is None else at.as_tensor_variable(lower)
+        upper = at.constant(np.inf) if upper is None else at.as_tensor_variable(upper)
+
+        # When size is not specified, dist may have to be broadcasted according to lower/upper
+        dist_shape = size if size is not None else at.broadcast_shape(dist, lower, upper)
+        dist = change_rv_size(dist, dist_shape)
 
         # Censoring is achieved by clipping the base distribution between lower and upper
         rv_out = at.clip(dist, lower, upper)
@@ -88,8 +92,6 @@ def rv_op(cls, dist, lower=None, upper=None, size=None, rngs=None):
         rv_out.tag.lower = lower
         rv_out.tag.upper = upper
 
-        if size is not None:
-            rv_out = cls.change_size(rv_out, size)
         if rngs is not None:
             rv_out = cls.change_rngs(rv_out, rngs)
 
@@ -101,12 +103,10 @@ def ndim_supp(cls, *dist_params):
 
     @classmethod
     def change_size(cls, rv, new_size, expand=False):
-        dist_node = rv.tag.dist.owner
+        dist = rv.tag.dist
         lower = rv.tag.lower
         upper = rv.tag.upper
-        rng, old_size, dtype, *dist_params = dist_node.inputs
-        new_size = new_size if not expand else tuple(new_size) + tuple(old_size)
-        new_dist = dist_node.op.make_node(rng, new_size, dtype, *dist_params).default_output()
+        new_dist = change_rv_size(dist, new_size, expand=expand)
         return cls.rv_op(new_dist, lower, upper)
 
     @classmethod

pymc/tests/test_distributions.py

@@ -3381,6 +3381,18 @@ def test_change_size(self):
         new_dist = pm.Censored.change_size(base_dist, (4,), expand=True)
         assert new_dist.eval().shape == (4, 3, 2)
 
+    def test_dist_broadcasted_by_lower_upper(self):
+        x = pm.Censored.dist(pm.Normal.dist(), lower=np.zeros((2,)), upper=None)
+        assert tuple(x.owner.inputs[0].shape.eval()) == (2,)
+
+        x = pm.Censored.dist(pm.Normal.dist(), lower=np.zeros((2,)), upper=np.zeros((4, 2)))
+        assert tuple(x.owner.inputs[0].shape.eval()) == (4, 2)
+
+        x = pm.Censored.dist(
+            pm.Normal.dist(size=(3, 4, 2)), lower=np.zeros((2,)), upper=np.zeros((4, 2))
+        )
+        assert tuple(x.owner.inputs[0].shape.eval()) == (3, 4, 2)
+
 
 class TestLKJCholeskCov:
     def test_dist(self):