Extra changes to fix some more tests

Author: lucianopaz

homepy/pytensorf.py

@@ -18,6 +18,7 @@
 import pytensor
 import pytensor.tensor as pt
 from pymc.logprob.utils import get_underlying_scalar_constant_value
+from pymc.pytensorf import find_rng_nodes
 from pytensor.graph import FunctionGraph
 from pytensor.graph.basic import clone_get_equiv
 from pytensor.tensor.random.basic import NormalRV
@@ -30,31 +31,6 @@
     from pytensor.graph.opt import local_optimizer as node_rewriter
 
 
-def find_rng_nodes(outputs):
-    """
-    Find all RNG (random number generator) shared variables in the graph.
-
-    In PyMC v5/PyTensor, this replaces pymc.pytensorf.find_rng_nodes from PyMC v4.
-    """
-    from pytensor.graph import graph_inputs
-    from pytensor.tensor.random.var import (
-        RandomGeneratorSharedVariable,
-        RandomStateSharedVariable,
-    )
-
-    # Get all inputs to the graph
-    inputs = graph_inputs(outputs)
-
-    # Filter for RNG shared variables
-    rng_nodes = [
-        node
-        for node in inputs
-        if isinstance(node, (RandomStateSharedVariable, RandomGeneratorSharedVariable))
-    ]
-
-    return rng_nodes
-
-
 def resampled_as_non_centered(outputs, resampled_vars, free_RVs):
     # We make a shallow copy of the lists so that we don't make alter the original ones
     resampled_vars = [[rv for rv in rvs] for rvs in resampled_vars]
@@ -95,7 +71,6 @@ def make_normal_not_centered(fgraph, node):
             return [None, new_node]
 
     rewrite = in2out(make_normal_not_centered)
-
     graph = FunctionGraph(outputs=outputs, clone=False)
     rewrite.apply(graph)
     return graph.outputs, free_RVs, resampled_vars
@@ -114,15 +89,9 @@ def clone_replace_rv_consistent(outputs, free_RVs, replace):
     # That way, the draws across the cloned and uncloned graph will be uncorrelated
     rng_nodes = find_rng_nodes(fg.outputs)
     new_rng_nodes: List[Union[np.random.RandomState, np.random.Generator]] = []
-    from pytensor.tensor.random.var import RandomStateSharedVariable
 
-    for rng_node in rng_nodes:
-        rng_cls: type
-        if isinstance(rng_node, RandomStateSharedVariable):
-            rng_cls = np.random.RandomState
-        else:
-            rng_cls = np.random.Generator
-        new_rng_nodes.append(pytensor.shared(rng_cls(np.random.PCG64())))
+    for _ in rng_nodes:
+        new_rng_nodes.append(pytensor.shared(np.random.Generator(np.random.PCG64())))
     orig_replace = {clone_map[rv]: rv for rv in free_RVs if rv in clone_map}
     orig_replace.update(dict(zip(rng_nodes, new_rng_nodes)))
     # replace_var can only be constant values or shared, not graph that depend on nodes

homepy/tests/models/test_base.py

@@ -852,6 +852,8 @@ def test_get_fixed_effect_observations(
         gdf = data.groupby(columns)
         assert len(gdf.groups) == len(fixed_effect_observations)
         for group, subdf in gdf:
+            if len(group) == 1:
+                group = group[0]
             assert np.array_equal(
                 subdf.loc[fixed_effect_observations[group], "bins"], consistent_time_bins
             )