Allow check_dist_not_registered to be called inside CustomDist

Author: ricardoV94

pymc/util.py

@@ -27,6 +27,8 @@
 from pytensor.compile import SharedVariable
 from pytensor.graph.utils import ValidatingScratchpad
 
+from pymc.exceptions import BlockModelAccessError
+
 
 class _UnsetType:
     """Type for the `UNSET` object to make it look nice in `help(...)` outputs."""
@@ -367,7 +369,7 @@ def check_dist_not_registered(dist, model=None):
 
     try:
         model = modelcontext(None)
-    except TypeError:
+    except (TypeError, BlockModelAccessError):
         pass
     else:
         if dist in model.basic_RVs:

tests/distributions/test_distribution.py

@@ -590,6 +590,13 @@ def shifted_normal(mu, sigma, size):
             pm.logp(pm.NormalMixture.dist(w=w, mu=mus, sigma=sds), test_value).eval(),
         )
 
+    def test_symbolic_dist(self):
+        # Test we can create a SymbolicDist inside a CustomDist
+        def dist(size):
+            return pm.Truncated.dist(pm.Beta.dist(1, 1, size=size), lower=0.1, upper=0.9)
+
+        assert pm.CustomDist.dist(dist=dist)
+
 
 class TestSymbolicRandomVariable:
     def test_inline(self):