Explicit rng in expand_dist_dims

Author: ricardoV94

pymc/dims/distributions/core.py

@@ -371,7 +371,8 @@ def expand_dist_dims(dist: XTensorVariable, extra_dims: dict[str, Any]) -> XTens
             dist_props["extra_dims"] = (*(extra_dims.keys()), *dist_props["extra_dims"])
             new_dist_op = type(dist.owner.op)(**dist_props)
             _old_rng, *params_and_dim_lengths = dist.owner.inputs
-            new_rng = None  # We don't propagate the old RNG, because we don't want the new and old dists to be correlated
+            # We don't propagate the old RNG, because we don't want the new and old dists to be correlated
+            new_rng = pt.random.shared_rng(seed=None)
             return new_dist_op(new_rng, *extra_dims.values(), *params_and_dim_lengths)
         case Transpose():
             return expand_dist_dims(dist.owner.inputs[0], extra_dims=extra_dims).transpose(

tests/dims/distributions/test_censored.py

@@ -24,6 +24,8 @@
 from pymc.model.core import Model
 from tests.dims.utils import assert_equivalent_logp_graph, assert_equivalent_random_graph
 
+pytestmark = pytest.mark.filterwarnings("error")
+
 
 @pytest.mark.parametrize("lower", [None, -1])
 @pytest.mark.parametrize("upper", [None, 1])

tests/dims/distributions/test_scalar.py

@@ -39,6 +39,8 @@
 )
 from tests.dims.utils import assert_equivalent_logp_graph, assert_equivalent_random_graph
 
+pytestmark = pytest.mark.filterwarnings("error")
+
 
 def test_flat():
     coords = {"a": range(3)}

tests/dims/distributions/test_vector.py

@@ -13,6 +13,7 @@
 #   limitations under the License.
 import numpy as np
 import pytensor.tensor as pt
+import pytest
 
 from pytensor.xtensor import as_xtensor
 
@@ -22,6 +23,8 @@
 from pymc.dims import Categorical, Dirichlet, MvNormal, ZeroSumNormal
 from tests.dims.utils import assert_equivalent_logp_graph, assert_equivalent_random_graph
 
+pytestmark = pytest.mark.filterwarnings("error")
+
 
 def test_categorical():
     coords = {"a": range(3), "b": range(4)}