Update distributions to handle new rng semantics

Author: jessegrabowski

pymc_extras/distributions/continuous.py

@@ -270,7 +270,10 @@ class Chi:
 
     @staticmethod
     def chi_dist(nu: TensorVariable, size: TensorVariable) -> TensorVariable:
-        return pt.math.sqrt(ChiSquared.dist(nu=nu, size=size))
+        _, rv = ChiSquared.dist(
+            nu=nu, size=size, rng=pt.random.shared_rng(seed=0), return_next_rng=True
+        )
+        return pt.math.sqrt(rv)
 
     def __new__(cls, name, nu, **kwargs):
         if "observed" not in kwargs:
@@ -331,7 +334,8 @@ def maxwell_dist(a: TensorVariable, size: TensorVariable) -> TensorVariable:
 
         a = CheckParameterValue("a > 0")(a, pt.all(pt.gt(a, 0)))
 
-        return Chi.dist(nu=3, size=size) * a
+        _, chi = Chi.dist(nu=3, size=size, rng=pt.random.shared_rng(seed=0), return_next_rng=True)
+        return chi * a
 
     def __new__(cls, name, a, **kwargs):
         return CustomDist(

pymc_extras/distributions/discrete.py

@@ -252,8 +252,13 @@ def beta_negative_binomial_dist(alpha, beta, r, size):
         if rv_size_is_none(size):
             alpha, beta, r = pt.broadcast_arrays(alpha, beta, r)
 
-        p = pm.Beta.dist(alpha, beta, size=size)
-        return pm.NegativeBinomial.dist(p, r, size=size)
+        _, p = pm.Beta.dist(
+            alpha, beta, size=size, rng=pt.random.shared_rng(seed=0), return_next_rng=True
+        )
+        _, rv = pm.NegativeBinomial.dist(
+            p, r, size=size, rng=pt.random.shared_rng(seed=0), return_next_rng=True
+        )
+        return rv
 
     @staticmethod
     def beta_negative_binomial_logp(value, alpha, beta, r):
@@ -361,7 +366,13 @@ def skellam_dist(mu1, mu2, size):
         if rv_size_is_none(size):
             mu1, mu2 = pt.broadcast_arrays(mu1, mu2)
 
-        return pm.Poisson.dist(mu=mu1, size=size) - pm.Poisson.dist(mu=mu2, size=size)
+        _, p1 = pm.Poisson.dist(
+            mu=mu1, size=size, rng=pt.random.shared_rng(seed=0), return_next_rng=True
+        )
+        _, p2 = pm.Poisson.dist(
+            mu=mu2, size=size, rng=pt.random.shared_rng(seed=0), return_next_rng=True
+        )
+        return p1 - p2
 
     @staticmethod
     def skellam_logp(value, mu1, mu2):

pymc_extras/distributions/timeseries.py

@@ -175,7 +175,9 @@ def dist(cls, P=None, logit_P=None, steps=None, init_dist=None, n_lags=1, **kwar
                 UserWarning,
             )
             k = P.shape[-1]
-            init_dist = pm.Categorical.dist(p=pt.full((k,), 1 / k))
+            _, init_dist = pm.Categorical.dist(
+                p=pt.full((k,), 1 / k), rng=pt.random.shared_rng(seed=0), return_next_rng=True
+            )
 
         return super().dist([P, steps, init_dist], n_lags=n_lags, **kwargs)
 
@@ -198,7 +200,7 @@ def rv_op(cls, P, steps, init_dist, n_lags, size=None):
         def transition(*args):
             old_rng, *states, transition_probs = args
             p = transition_probs[tuple(states)]
-            next_rng, next_state = pm.Categorical.dist(p=p, rng=old_rng).owner.outputs
+            next_rng, next_state = pm.Categorical.dist(p=p, rng=old_rng, return_next_rng=True)
             return next_rng, next_state
 
         state_next_rng, markov_chain = pytensor.scan(

pymc_extras/inference/advi/autoguide.py

@@ -87,7 +87,14 @@ def AutoDiagonalNormal(model: Model) -> AutoGuideModel:
             loc = pt.tensor(f"{rv.name}_loc", shape=rv.type.shape)
             scale = pt.tensor(f"{rv.name}_scale", shape=rv.type.shape)
             # TODO: Make these customizable
-            params_init_values[loc] = pt.random.uniform(-1, 1, size=free_rv_shapes[rv]).eval()
+            _, loc_init = pt.random.uniform(
+                -1,
+                1,
+                size=free_rv_shapes[rv],
+                rng=pt.random.shared_rng(seed=0),
+                return_next_rng=True,
+            )
+            params_init_values[loc] = loc_init.eval()
             params_init_values[scale] = pt.full(free_rv_shapes[rv], 0.1).eval()
 
             z = Normal(

pymc_extras/inference/laplace_approx/laplace.py

@@ -232,10 +232,23 @@ def draws_from_laplace_approx(
     size = (draws,) if vectorize_draws else ()
     if covariance is not None:
         sigma_pt = pt.matrix("cov", shape=(n, n), dtype=covariance.dtype)
-        laplace_approximation = pm.MvNormal.dist(mu=mu_pt, cov=sigma_pt, size=size, method="svd")
+        _, laplace_approximation = pm.MvNormal.dist(
+            mu=mu_pt,
+            cov=sigma_pt,
+            size=size,
+            method="svd",
+            rng=pt.random.shared_rng(seed=0),
+            return_next_rng=True,
+        )
     else:
         sigma_pt = pt.vector("sigma", shape=(n,), dtype=standard_deviation.dtype)
-        laplace_approximation = pm.Normal.dist(mu=mu_pt, sigma=sigma_pt, size=(*size, n))
+        _, laplace_approximation = pm.Normal.dist(
+            mu=mu_pt,
+            sigma=sigma_pt,
+            size=(*size, n),
+            rng=pt.random.shared_rng(seed=0),
+            return_next_rng=True,
+        )
 
     constrained_vars = unpack_last_axis(
         laplace_approximation,

pymc_extras/model/transforms/autoreparam.py

@@ -231,11 +231,12 @@ def _(
     rng, size, loc, scale = node.inputs
     if transform is not None:
         raise NotImplementedError("Reparametrization of Normal with Transform is not implemented")
-    vip_rv_ = pm.Normal.dist(
+    _, vip_rv_ = pm.Normal.dist(
         lam * loc,
         scale**lam,
         size=size,
         rng=rng,
+        return_next_rng=True,
     )
     vip_rv_.name = f"{name}::tau_"
 
@@ -266,10 +267,11 @@ def _(
     rng, size, scale = node.inputs
     scale_centered = scale**lam
     scale_noncentered = scale ** (1 - lam)
-    vip_rv_ = pm.Exponential.dist(
+    _, vip_rv_ = pm.Exponential.dist(
         scale=scale_centered,
         size=size,
         rng=rng,
+        return_next_rng=True,
     )
     vip_rv_value_ = vip_rv_.clone()
     vip_rv_.name = f"{name}::tau_"

pymc_extras/statespace/filters/distributions.py

@@ -180,8 +180,6 @@ def rv_op(
         ]
         non_sequences = [x for x in [c_, d_, T_, Z_, R_, H_, Q_] if x not in sequences]
 
-        rng = normalize_rng_param(rng)
-
         def sort_args(args):
             sorted_args = []
 
@@ -206,11 +204,11 @@ def step_fn(*args):
             a = state[:k]
 
             middle_rng, a_innovation = pm.MvNormal.dist(
-                mu=0, cov=Q, rng=rng, method=method
-            ).owner.outputs
+                mu=0, cov=Q, rng=rng, method=method, return_next_rng=True
+            )
             next_rng, y_innovation = pm.MvNormal.dist(
-                mu=0, cov=H, rng=middle_rng, method=method
-            ).owner.outputs
+                mu=0, cov=H, rng=middle_rng, method=method, return_next_rng=True
+            )
 
             a_mu = c + T @ a
             a_next = a_mu + R @ a_innovation
@@ -225,14 +223,18 @@ def step_fn(*args):
         Z_init = Z_ if Z_ in non_sequences else Z_[0]
         H_init = H_ if H_ in non_sequences else H_[0]
 
-        init_x_ = pm.MvNormal.dist(a0_, P0_, rng=rng, method=method)
-        init_y_ = pm.MvNormal.dist(Z_init @ init_x_, H_init, rng=rng, method=method)
+        rng = normalize_rng_param(rng)
+
+        next_rng, init_x_ = pm.MvNormal.dist(a0_, P0_, rng=rng, method=method, return_next_rng=True)
+        next_rng, init_y_ = pm.MvNormal.dist(
+            Z_init @ init_x_, H_init, rng=next_rng, method=method, return_next_rng=True
+        )
 
         init_dist_ = pt.concatenate([init_x_, init_y_], axis=0)
 
         ss_rng, statespace = pytensor.scan(
             step_fn,
-            outputs_info=[rng, init_dist_],
+            outputs_info=[next_rng, init_dist_],
             sequences=None if len(sequences) == 0 else sequences,
             non_sequences=[*non_sequences],
             n_steps=steps,
@@ -379,8 +381,8 @@ def rv_op(cls, mus, covs, logp, method="svd", size=None, rng=None):
 
         mus_, covs_ = mus.type(), covs.type()
         seq_mvn_rng, mvn_seq = multivariate_normal(
-            mean=mus_, cov=covs_, rng=rng, method=method
-        ).owner.outputs
+            mean=mus_, cov=covs_, rng=rng, method=method, return_next_rng=True
+        )
 
         mvn_seq_op = KalmanFilterRV(
             inputs=[mus_, covs_, logp_, rng], outputs=[seq_mvn_rng, mvn_seq], ndim_supp=2

tests/model/marginal/test_graph_analysis.py

@@ -14,11 +14,15 @@
 def test_is_conditional_dependent_static_shape():
     """Test that we don't consider dependencies through "constant" shape Ops"""
     x1 = pt.matrix("x1", shape=(None, 5))
-    y1 = pt.random.normal(size=pt.shape(x1))
+    _, y1 = pt.random.normal(
+        size=pt.shape(x1), rng=pt.random.shared_rng(seed=0), return_next_rng=True
+    )
     assert is_conditional_dependent(y1, x1, [x1, y1])
 
     x2 = pt.matrix("x2", shape=(9, 5))
-    y2 = pt.random.normal(size=pt.shape(x2))
+    _, y2 = pt.random.normal(
+        size=pt.shape(x2), rng=pt.random.shared_rng(seed=0), return_next_rng=True
+    )
     assert not is_conditional_dependent(y2, x2, [x2, y2])
 
 
@@ -145,25 +149,36 @@ def test_blockwise(self):
     def test_random_variable(self):
         inp = pt.tensor(shape=(5, 4, 3))
 
-        out1 = pt.random.normal(loc=inp)
-        out2 = pt.random.categorical(p=inp[..., None])
-        out3 = pt.random.multivariate_normal(mean=inp[..., None], cov=pt.eye(1))
+        _, out1 = pt.random.normal(loc=inp, rng=pt.random.shared_rng(seed=0), return_next_rng=True)
+        _, out2 = pt.random.categorical(
+            p=inp[..., None], rng=pt.random.shared_rng(seed=0), return_next_rng=True
+        )
+        _, out3 = pt.random.multivariate_normal(
+            mean=inp[..., None],
+            cov=pt.eye(1),
+            rng=pt.random.shared_rng(seed=0),
+            return_next_rng=True,
+        )
         [dims1, dims2, dims3] = subgraph_batch_dim_connection(inp, [out1, out2, out3])
         assert dims1 == (0, 1, 2)
         assert dims2 == (0, 1, 2)
         assert dims3 == (0, 1, 2, None)
 
-        invalid_out = pt.random.categorical(p=inp)
+        _, invalid_out = pt.random.categorical(
+            p=inp, rng=pt.random.shared_rng(seed=0), return_next_rng=True
+        )
         with pytest.raises(ValueError, match="Use of known dimensions"):
             subgraph_batch_dim_connection(inp, [invalid_out])
 
-        invalid_out = pt.random.multivariate_normal(mean=inp, cov=pt.eye(3))
+        _, invalid_out = pt.random.multivariate_normal(
+            mean=inp, cov=pt.eye(3), rng=pt.random.shared_rng(seed=0), return_next_rng=True
+        )
         with pytest.raises(ValueError, match="Use of known dimensions"):
             subgraph_batch_dim_connection(inp, [invalid_out])
 
     def test_minibatched_random_variable(self):
         inp = pt.tensor(shape=(4, 3, 2))
-        out1 = pt.random.normal(loc=inp)
+        _, out1 = pt.random.normal(loc=inp, rng=pt.random.shared_rng(seed=0), return_next_rng=True)
         out2 = create_minibatch_rv(out1, total_size=(10, 10, 10))
         [dims1] = subgraph_batch_dim_connection(inp, [out2])
         assert dims1 == (0, 1, 2)
@@ -174,7 +189,9 @@ def test_symbolic_random_variable(self):
         # Test univariate
         out = CustomDist.dist(
             inp,
-            dist=lambda mu, size: pt.random.normal(loc=mu, size=size),
+            dist=lambda mu, size: pt.random.normal(
+                loc=mu, size=size, rng=pt.random.shared_rng(seed=0), return_next_rng=True
+            )[1],
         )
         [dims] = subgraph_batch_dim_connection(inp, [out])
         assert dims == (0, 1, 2)
@@ -183,7 +200,13 @@ def test_symbolic_random_variable(self):
         def dist(mu, size):
             if isinstance(size.type, NoneTypeT):
                 size = mu.shape
-            return pt.random.normal(loc=mu[..., None], size=(*size, 2))
+            _, rv = pt.random.normal(
+                loc=mu[..., None],
+                size=(*size, 2),
+                rng=pt.random.shared_rng(seed=0),
+                return_next_rng=True,
+            )
+            return rv
 
         out = CustomDist.dist(inp, dist=dist, size=(4, 3, 2), signature="()->(2)")
         [dims] = subgraph_batch_dim_connection(inp, [out])