Update SymbolicRandomVariables to manage RNGs explicitly

* Also fixes error when jaxifying MarginalModel logp

Author: ricardoV94

pymc_experimental/distributions/timeseries.py

@@ -202,13 +202,13 @@ def transition(*args):
         discrete_mc_ = pt.moveaxis(pt.concatenate([init_dist_, markov_chain], axis=0), 0, -1)
 
         discrete_mc_op = DiscreteMarkovChainRV(
-            inputs=[P_, steps_, init_dist_],
+            inputs=[P_, steps_, init_dist_, state_rng],
             outputs=[state_next_rng, discrete_mc_],
             ndim_supp=1,
             n_lags=n_lags,
         )
 
-        discrete_mc = discrete_mc_op(P, steps, init_dist)
+        discrete_mc = discrete_mc_op(P, steps, init_dist, state_rng)
         return discrete_mc
 
 

pymc_experimental/model/marginal_model.py

@@ -410,7 +410,7 @@ def transform_input(inputs):
                 marginalized_rv.type, dependent_logps
             )
 
-            rv_shape = constant_fold(tuple(marginalized_rv.shape))
+            rv_shape = constant_fold(tuple(marginalized_rv.shape), raise_not_constant=False)
             rv_domain = get_domain_of_finite_discrete_rv(marginalized_rv)
             rv_domain_tensor = pt.moveaxis(
                 pt.full(
@@ -579,6 +579,15 @@ def is_elemwise_subgraph(rv_to_marginalize, other_input_rvs, output_rvs):
     return True
 
 
+from pytensor.graph.basic import graph_inputs
+
+
+def collect_shared_vars(outputs, blockers):
+    return [
+        inp for inp in graph_inputs(outputs, blockers=blockers) if isinstance(inp, SharedVariable)
+    ]
+
+
 def replace_finite_discrete_marginal_subgraph(fgraph, rv_to_marginalize, all_rvs):
     # TODO: This should eventually be integrated in a more general routine that can
     #  identify other types of supported marginalization, of which finite discrete
@@ -621,27 +630,21 @@ def replace_finite_discrete_marginal_subgraph(fgraph, rv_to_marginalize, all_rvs
     rvs_to_marginalize = [rv_to_marginalize, *dependent_rvs]
 
     outputs = rvs_to_marginalize
-    # Clone replace inner RV rng inputs so that we can be sure of the update order
-    # replace_inputs = {rng: rng.type() for rng in updates_rvs_to_marginalize.keys()}
-    # Clone replace outter RV inputs, so that their shared RNGs don't make it into
-    # the inner graph of the marginalized RVs
-    # FIXME: This shouldn't be needed!
-    replace_inputs = {}
-    replace_inputs.update({input_rv: input_rv.type() for input_rv in input_rvs})
-    cloned_outputs = clone_replace(outputs, replace=replace_inputs)
+    # We are strict about shared variables in SymbolicRandomVariables
+    inputs = input_rvs + collect_shared_vars(rvs_to_marginalize, blockers=input_rvs)
 
     if isinstance(rv_to_marginalize.owner.op, DiscreteMarkovChain):
         marginalize_constructor = DiscreteMarginalMarkovChainRV
     else:
         marginalize_constructor = FiniteDiscreteMarginalRV
 
     marginalization_op = marginalize_constructor(
-        inputs=list(replace_inputs.values()),
-        outputs=cloned_outputs,
+        inputs=inputs,
+        outputs=outputs,
         ndim_supp=ndim_supp,
     )
 
-    marginalized_rvs = marginalization_op(*replace_inputs.keys())
+    marginalized_rvs = marginalization_op(*inputs)
     fgraph.replace_all(tuple(zip(rvs_to_marginalize, marginalized_rvs)))
     return rvs_to_marginalize, marginalized_rvs
 

pymc_experimental/statespace/filters/distributions.py

@@ -193,12 +193,12 @@ def step_fn(*args):
 
         (ss_rng,) = tuple(updates.values())
         linear_gaussian_ss_op = LinearGaussianStateSpaceRV(
-            inputs=[a0_, P0_, c_, d_, T_, Z_, R_, H_, Q_, steps_],
+            inputs=[a0_, P0_, c_, d_, T_, Z_, R_, H_, Q_, steps_, rng],
             outputs=[ss_rng, statespace_],
             ndim_supp=1,
         )
 
-        linear_gaussian_ss = linear_gaussian_ss_op(a0, P0, c, d, T, Z, R, H, Q, steps)
+        linear_gaussian_ss = linear_gaussian_ss_op(a0, P0, c, d, T, Z, R, H, Q, steps, rng)
         return linear_gaussian_ss
 
 
@@ -354,10 +354,10 @@ def step(mu, cov, rng):
         (seq_mvn_rng,) = tuple(updates.values())
 
         mvn_seq_op = KalmanFilterRV(
-            inputs=[mus_, covs_, logp_, steps_], outputs=[seq_mvn_rng, mvn_seq], ndim_supp=2
+            inputs=[mus_, covs_, logp_, steps_, rng], outputs=[seq_mvn_rng, mvn_seq], ndim_supp=2
         )
 
-        mvn_seq = mvn_seq_op(mus, covs, logp, steps)
+        mvn_seq = mvn_seq_op(mus, covs, logp, steps, rng)
         return mvn_seq
 
 

pymc_experimental/tests/model/test_marginal_model.py

@@ -60,9 +60,9 @@ def test_marginalized_bernoulli_logp():
         [idx, y],
         ndim_supp=0,
         n_updates=0,
-    )(
-        mu
-    )[0].owner
+        # Ignore the fact we didn't specify shared RNG input/outputs for idx,y
+        strict=False,
+    )(mu)[0].owner
 
     y_vv = y.clone()
     (logp,) = _logprob(
@@ -758,3 +758,19 @@ def test_marginalized_hmm_multiple_emissions(batch_emission1, batch_emission2):
     test_value_emission2 = np.broadcast_to(-test_value, emission2_shape)
     test_point = {"emission_1": test_value_emission1, "emission_2": test_value_emission2}
     np.testing.assert_allclose(logp_fn(test_point), expected_logp)
+
+
+def test_mutable_indexing_jax_backend():
+    pytest.importorskip("jax")
+    from pymc.sampling.jax import get_jaxified_logp
+
+    with MarginalModel() as model:
+        data = pm.Data(f"data", np.zeros(10))
+
+        cat_effect = pm.Normal("cat_effect", sigma=1, shape=5)
+        cat_effect_idx = pm.Data("cat_effect_idx", np.array([0, 1] * 5))
+
+        is_outlier = pm.Bernoulli("is_outlier", 0.4, shape=10)
+        pm.LogNormal("y", mu=cat_effect[cat_effect_idx], sigma=1 + is_outlier, observed=data)
+    model.marginalize(["is_outlier"])
+    get_jaxified_logp(model)