Raise error when trying to jaxify graphs with RNGs

Author: ricardoV94

pymc/sampling/jax.py

@@ -18,6 +18,8 @@
 from functools import partial
 from typing import Any, Callable, Dict, List, Optional, Sequence, Union
 
+from pytensor.tensor.random.type import RandomType
+
 from pymc.initial_point import StartDict
 from pymc.sampling.mcmc import _init_jitter
 
@@ -81,6 +83,11 @@ def _replace_shared_variables(graph: List[TensorVariable]) -> List[TensorVariabl
 
     shared_variables = [var for var in graph_inputs(graph) if isinstance(var, SharedVariable)]
 
+    if any(isinstance(var.type, RandomType) for var in shared_variables):
+        raise ValueError(
+            "Graph contains shared RandomType variables which cannot be safely replaced"
+        )
+
     if any(var.default_update is not None for var in shared_variables):
         raise ValueError(
             "Graph contains shared variables with default_update which cannot "

pymc/tests/sampling/test_jax.py

@@ -29,6 +29,18 @@
 
 import pymc as pm
 
+from pymc.sampling.jax import (
+    _get_batched_jittered_initial_points,
+    _get_log_likelihood,
+    _numpyro_nuts_defaults,
+    _replace_shared_variables,
+    _update_numpyro_nuts_kwargs,
+    get_jaxified_graph,
+    get_jaxified_logp,
+    sample_blackjax_nuts,
+    sample_numpyro_nuts,
+)
+
 
 def test_old_import_route():
     import pymc.sampling.jax as new_sj
@@ -37,20 +49,6 @@ def test_old_import_route():
     assert set(new_sj.__all__) <= set(dir(old_sj))
 
 
-with pytest.warns(UserWarning, match="module is experimental"):
-    from pymc.sampling.jax import (
-        _get_batched_jittered_initial_points,
-        _get_log_likelihood,
-        _numpyro_nuts_defaults,
-        _replace_shared_variables,
-        _update_numpyro_nuts_kwargs,
-        get_jaxified_graph,
-        get_jaxified_logp,
-        sample_blackjax_nuts,
-        sample_numpyro_nuts,
-    )
-
-
 @pytest.mark.parametrize(
     "sampler",
     [
@@ -182,6 +180,11 @@ def test_replace_shared_variables():
     with pytest.raises(ValueError, match="shared variables with default_update"):
         _replace_shared_variables([x])
 
+    shared_rng = pytensor.shared(np.random.default_rng(), name="shared_rng")
+    x = pytensor.tensor.random.normal(rng=shared_rng)
+    with pytest.raises(ValueError, match="Graph contains shared RandomType variables"):
+        _replace_shared_variables([x])
+
 
 def test_get_jaxified_logp():
     with pm.Model() as m: