Allow explicit rng for CustomDist that only require one

Author: ricardoV94

pymc/distributions/custom.py

@@ -259,6 +259,7 @@ def rv_op(
         size=None,
         signature: str,
         class_name: str,
+        rng=None,
     ):
         size = normalize_size_param(size)
         # If it's NoneConst, just use that as the dummy
@@ -270,7 +271,8 @@ def rv_op(
             dummy_rv = dist(*dummy_dist_params, dummy_size_param)
         dummy_params = [dummy_size_param, *dummy_dist_params]
         # RNGs are not passed as explicit inputs (because we usually don't know how many are needed)
-        # We retrieve them here. This will also raise if the user forgot to specify some update in a Scan Op
+        # We retrieve them here. This will also raise if the user forgot to specify some update in an InnerGraphOp (e.g., Scan)
+        # If the user passed an explicit rng we will respect that later when we instantiate the final rv_op
         dummy_updates_dict = collect_default_updates(inputs=dummy_params, outputs=(dummy_rv,))
 
         rv_type = type(
@@ -357,6 +359,14 @@ def change_custom_dist_size(op, rv, new_size, expand):
             outputs=outputs,
             extended_signature=extended_signature,
         )
+        if rng is not None:
+            # User passed an RNG, use that if the graph only required one, raise otherwise
+            if len(rngs) != 1:
+                raise ValueError(
+                    f"CustomDist received an explicit rng but it actually requires {len(rngs)} rngs."
+                    " Please modify your dist function to only use one rng, or don't pass an explicitly rng."
+                )
+            rngs = (rng,)
         return rv_op(size, *dist_params, *rngs)
 
     @staticmethod

tests/distributions/test_custom.py

@@ -708,3 +708,47 @@ def normal_shifted(mu, size):
             observed_logp.eval({latent_vv: latent_vv_test, observed_vv: observed_vv_test}),
             expected_logp,
         )
+
+    def test_explicit_rng(self):
+        def custom_dist(mu, size):
+            return Normal.dist(mu, size=size)
+
+        x = CustomDist.dist(0, dist=custom_dist)
+        assert len(x.owner.op.rng_params(x.owner)) == 1  # Rng created by default
+
+        explicit_rng = pt.random.type.random_generator_type("rng")
+        x_explicit = CustomDist.dist(0, dist=custom_dist, rng=explicit_rng)
+        [used_rng] = x_explicit.owner.op.rng_params(x_explicit.owner)
+        assert used_rng is explicit_rng
+
+        # API for passing multiple explicit RNGs not supported
+        def custom_dist_multi_rng(mu, size):
+            return Normal.dist(mu, size=size) + Normal.dist(0, size=size)
+
+        x = CustomDist.dist(0, dist=custom_dist_multi_rng)
+        assert len(x.owner.op.rng_params(x.owner)) == 2
+
+        with pytest.raises(
+            ValueError,
+            match="CustomDist received an explicit rng but it actually requires 2 rngs",
+        ):
+            CustomDist.dist(
+                0,
+                dist=custom_dist_multi_rng,
+                rng=explicit_rng,
+            )
+
+        # But it can be done if the custom_dist uses only one RNG internally
+        def custom_dist_multi_rng_fixed(mu, size):
+            next_rng, x = Normal.dist(mu, size=size).owner.outputs
+            return x + Normal.dist(0, size=size, rng=next_rng)
+
+        x = CustomDist.dist(0, dist=custom_dist_multi_rng_fixed)
+        assert len(x.owner.op.rng_params(x.owner)) == 1
+        x_explicit = CustomDist.dist(
+            0,
+            dist=custom_dist_multi_rng_fixed,
+            rng=explicit_rng,
+        )
+        [used_rng] = x_explicit.owner.op.rng_params(x_explicit.owner)
+        assert used_rng is explicit_rng