From 8804bcac3de3bb7f96bf7c453fe6be80302c1b8c Mon Sep 17 00:00:00 2001 From: nileshpatil6 Date: Wed, 3 Jun 2026 22:51:47 +0530 Subject: [PATCH] fix(infer): render_model now draws edges through deterministic sites pyro.render_model omitted edges when a sample site received a tensor from a pyro.deterministic site or a pyro.sample wrapped in a Delta transform. get_model_relations used TrackProvenance with include_deterministic passed through, so deterministic and Delta sites were not wrapped with their own provenance anchor. Downstream sites then resolved their parents transitively, skipping the intermediate deterministic/Delta nodes. Fix by always using TrackProvenance(include_deterministic=True) inside get_model_relations so every site acts as a provenance anchor. This correctly builds direct-parent edges regardless of the include_deterministic output filter. Update test_get_model_relations to reflect correct edges (i now lists f, g, h as direct parents instead of resolving them away to e). Fixes #3441 Signed-off-by: nileshpatil6 --- pyro/infer/inspect.py | 2 +- tests/infer/test_inspect.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/pyro/infer/inspect.py b/pyro/infer/inspect.py index 2580301dba..18055bba38 100644 --- a/pyro/infer/inspect.py +++ b/pyro/infer/inspect.py @@ -298,7 +298,7 @@ def model(data): assert isinstance(model_kwargs, dict) with torch.random.fork_rng(), torch.no_grad(), pyro.validation_enabled(False): - with TrackProvenance(include_deterministic=include_deterministic): + with TrackProvenance(include_deterministic=True): trace = poutine.trace(model).get_trace(*model_args, **model_kwargs) sample_sample = {} diff --git a/tests/infer/test_inspect.py b/tests/infer/test_inspect.py index 8b72c621ca..baccf00500 100644 --- a/tests/infer/test_inspect.py +++ b/tests/infer/test_inspect.py @@ -526,7 +526,7 @@ def model(data): "f": ["e"], "g": ["e"], "h": ["e"], - "i": ["e"], + "i": ["e", "f", "g", "h"], }, "sample_param": { "a": [],