Fix pyro parameter nodes not showing up in graph of pyro deterministi…

…c nodes [bugfix] (#3378) * Fixes pyro parameter nodes not showing up in graph of pyro deterministic nodes. * Linting and formatting. * Ignore existing object redefinitions as they are now checked by ruff. --------- Co-authored-by: Ben Zickel <[email protected]>
pyro-ppl · Jun 27, 2024 · 1d7da54 · 1d7da54
1 parent 64e71ee
commit 1d7da54
Show file tree

Hide file tree

Showing 4 changed files with 13 additions and 3 deletions.
diff --git a/pyro/distributions/hmm.py b/pyro/distributions/hmm.py
@@ -1091,7 +1091,7 @@ def __init__(
         self.transforms = transforms
 
     @constraints.dependent_property(event_dim=2)
-    def support(self):
+    def support(self):  # noqa: F811
         return constraints.independent(self.observation_dist.support, 1)
 
     def expand(self, batch_shape, _instance=None):

diff --git a/pyro/distributions/stable_log_prob.py b/pyro/distributions/stable_log_prob.py
@@ -44,7 +44,7 @@ def set_integrator(num_points):
 
 # Stub which is replaced by the default integrator when called for the first time
 # if a default integrator has not already been set.
-def integrate(*args, **kwargs):
+def integrate(*args, **kwargs):  # noqa: F811
     set_integrator(num_points=501)
     return integrate(*args, **kwargs)
 

diff --git a/pyro/infer/inspect.py b/pyro/infer/inspect.py
@@ -331,7 +331,7 @@ def _get_type_from_frozenname(frozen_name):
 
         sample_param[name] = [
             upstream
-            for upstream in get_provenance(site["fn"].log_prob(site["value"]))
+            for upstream in provenance
             if upstream != name and _get_type_from_frozenname(upstream) == "param"
         ]
 

diff --git a/tests/infer/test_util.py b/tests/infer/test_util.py
@@ -72,3 +72,13 @@ def guide(zdim=1, scale=1.0):
         scale=scale,
     )
     assert k > krange[0] and k < krange[1]
+
+
+def test_render_model_deterministic_param():
+    def model():
+        value = pyro.param("param", torch.tensor(0.0))
+        pyro.deterministic("deterministic", value)
+
+    graph = pyro.render_model(model, render_params=True, render_deterministic=True)
+
+    assert "\tparam -> deterministic\n" in graph.body