Add expand kwarg to Censored.change_size

Author: ricardoV94

pymc/distributions/censored.py

@@ -100,11 +100,12 @@ def ndim_supp(cls, *dist_params):
         return 0
 
     @classmethod
-    def change_size(cls, rv, new_size):
+    def change_size(cls, rv, new_size, expand=False):
         dist_node = rv.tag.dist.owner
         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()
         return cls.rv_op(new_dist, lower, upper)
 

pymc/tests/test_distributions.py

@@ -3357,6 +3357,15 @@ def test_censored_invalid_dist(self):
             ):
                 x = pm.Censored("x", registered_dist, lower=None, upper=None)
 
+    def test_change_size(self):
+        base_dist = pm.Censored.dist(pm.Normal.dist(), -1, 1, size=(3, 2))
+
+        new_dist = pm.Censored.change_size(base_dist, (4,))
+        assert new_dist.eval().shape == (4,)
+
+        new_dist = pm.Censored.change_size(base_dist, (4,), expand=True)
+        assert new_dist.eval().shape == (4, 3, 2)
+
 
 class TestLKJCholeskCov:
     def test_dist(self):