Ignore inner unused RNG inputs in collect_default_updates

Author: ricardoV94

pymc/pytensorf.py

@@ -855,35 +855,44 @@ def find_default_update(clients, rng: Variable) -> None | Variable:
 
         # Root case, RNG is not used elsewhere
         if not rng_clients:
-            return rng
+            return None
 
         if len(rng_clients) > 1:
             # Multiple clients are techincally fine if they are used in identical operations
             # We check if the default_update of each client would be the same
-            update, *other_updates = (
+            all_updates = [
                 find_default_update(
                     # Pass version of clients that includes only one the RNG clients at a time
                     clients | {rng: [rng_client]},
                     rng,
                 )
                 for rng_client in rng_clients
-            )
-            if all(equal_computations([update], [other_update]) for other_update in other_updates):
-                return update
-
-            warnings.warn(
-                f"RNG Variable {rng} has multiple distinct clients {rng_clients}, "
-                f"likely due to an inconsistent random graph. "
-                f"No default update will be returned.",
-                UserWarning,
-            )
-            return None
+            ]
+            updates = [update for update in all_updates if update is not None]
+            if not updates:
+                return None
+            if len(updates) == 1:
+                return updates[0]
+            else:
+                update, *other_updates = updates
+                if all(
+                    equal_computations([update], [other_update]) for other_update in other_updates
+                ):
+                    return update
+
+                warnings.warn(
+                    f"RNG Variable {rng} has multiple distinct clients {rng_clients}, "
+                    f"likely due to an inconsistent random graph. "
+                    f"No default update will be returned.",
+                    UserWarning,
+                )
+                return None
 
         [client, _] = rng_clients[0]
 
         # RNG is an output of the function, this is not a problem
         if isinstance(client.op, Output):
-            return rng
+            return None
 
         # RNG is used by another operator, which should output an update for the RNG
         if isinstance(client.op, RandomVariable):
@@ -912,18 +921,26 @@ def find_default_update(clients, rng: Variable) -> None | Variable:
                 )
         elif isinstance(client.op, OpFromGraph):
             try:
-                next_rng = collect_default_updates_inner_fgraph(client)[rng]
-            except (ValueError, KeyError):
+                next_rng = collect_default_updates_inner_fgraph(client).get(rng)
+                if next_rng is None:
+                    # OFG either does not make use of this RNG or inconsistent use that will have emitted a warning
+                    return None
+            except ValueError as exc:
                 raise ValueError(
                     f"No update found for at least one RNG used in OpFromGraph Op {client.op}.\n"
                     "You can use `pytensorf.collect_default_updates` and include those updates as outputs."
-                )
+                ) from exc
         else:
             # We don't know how this RNG should be updated. The user should provide an update manually
             return None
 
         # Recurse until we find final update for RNG
-        return find_default_update(clients, next_rng)
+        nested_next_rng = find_default_update(clients, next_rng)
+        if nested_next_rng is None:
+            # There were no more uses of this next_rng
+            return next_rng
+        else:
+            return nested_next_rng
 
     if inputs is None:
         inputs = []

tests/test_pytensorf.py

@@ -619,6 +619,20 @@ def test_op_from_graph_updates(self):
         fn = compile([], x, random_seed=1)
         assert not (set(fn()) & set(fn()))
 
+    def test_unused_ofg_rng(self):
+        rng = pytensor.shared(np.random.default_rng())
+        next_rng, x = pt.random.normal(rng=rng).owner.outputs
+        ofg1 = OpFromGraph([rng], [next_rng, x])
+        ofg2 = OpFromGraph([rng, x], [x + 1])
+
+        next_rng, x = ofg1(rng)
+        y = ofg2(rng, x)
+
+        # In all these cases the update should be the same
+        assert collect_default_updates([x]) == {rng: next_rng}
+        assert collect_default_updates([y]) == {rng: next_rng}
+        assert collect_default_updates([x, y]) == {rng: next_rng}
+
 
 def test_replace_rng_nodes():
     rng = pytensor.shared(np.random.default_rng())