fix file path when render model (#1857)

Author: fehiepsi

numpyro/infer/inspect.py

@@ -629,9 +629,14 @@ def render_model(
 
     if filename is not None:
         filename = Path(filename)
+        # remove leading period from suffix
+        filename_without_suffix = filename.with_suffix("")
         graph.render(
-            filename.stem, view=False, cleanup=True, format=filename.suffix[1:]
-        )  # remove leading period from suffix
+            filename_without_suffix,
+            view=False,
+            cleanup=True,
+            format=filename.suffix[1:],
+        )
 
     return graph
 

test/test_model_rendering.py

@@ -1,14 +1,20 @@
 # Copyright Contributors to the Pyro project.
 # SPDX-License-Identifier: Apache-2.0
 
+import os
+
 import numpy as np
 import pytest
 
 import jax.numpy as jnp
 
 import numpyro
 import numpyro.distributions as dist
-from numpyro.infer.inspect import generate_graph_specification, get_model_relations
+from numpyro.infer.inspect import (
+    generate_graph_specification,
+    get_model_relations,
+    render_model,
+)
 
 
 def simple(data):
@@ -129,3 +135,12 @@ def test_model_transformation(test_model, model_kwargs, expected_graph_spec):
     graph_spec = generate_graph_specification(relations)
 
     assert graph_spec == expected_graph_spec
+
+
+def test_render_model_filename():
+    def model():
+        numpyro.sample("x", dist.Normal(0, 1))
+
+    render_model(model, filename="graph.png")
+    assert os.path.exists("graph.png")
+    os.remove("graph.png")