We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
`numpyro.render_model` behaviour changed in 0.9.0
I noticed that the behaviour of `numpyro.render_model` appears to have changed since the latest release, specifically arrows signifying conditional dependence are not being rendered properly.

I haven't had time to dig into why this is yet, but here is a simple recreation

```python
import numpyro
import numpyro.distributions as dist


def model():
    theta = numpyro.sample("theta", dist.Uniform())
    
    with numpyro.plate("N", size=10):
        numpyro.sample("y", dist.Bernoulli(theta))
        
        
if __name__ == "__main__":
    g = numpyro.render_model(model)
    g.render(f"model-{numpyro.__version__}", format="png")
```

Running this with 0.8.0 and 0.9.0 I get (respectively)

![model-0 8 0](https://user-images.githubusercontent.com/15220906/152194532-07f70d39-8197-41c5-9e14-beefc33b807e.png)

![model-0 9 0](https://user-images.githubusercontent.com/15220906/152194539-98a41b9f-a673-4937-b864-1fb9e6c3d420.png)



--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `setup.py`
Content:
```
1 # Copyright Contributors to the Pyro project.
2 # SPDX-License-Identifier: Apache-2.0
3 
4 from __future__ import absolute_import, division, print_function
5 
6 import os
7 import sys
8 
9 from setuptools import find_packages, setup
10 
11 PROJECT_PATH = os.path.dirname(os.path.abspath(__file__))
12 _jax_version_constraints = ">=0.2.13,<0.2.28"
13 _jaxlib_version_constraints = ">=0.1.65"
14 
15 # Find version
16 for line in open(os.path.join(PROJECT_PATH, "numpyro", "version.py")):
17     if line.startswith("__version__ = "):
18         version = line.strip().split()[2][1:-1]
19 
20 # READ README.md for long description on PyPi.
21 try:
22     long_description = open("README.md", encoding="utf-8").read()
23 except Exception as e:
24     sys.stderr.write("Failed to read README.md:\n  {}\n".format(e))
25     sys.stderr.flush()
26     long_description = ""
27 
28 setup(
29     name="numpyro",
30     version=version,
31     description="Pyro PPL on NumPy",
32     packages=find_packages(include=["numpyro", "numpyro.*"]),
33     url="https://github.com/pyro-ppl/numpyro",
34     author="Uber AI Labs",
35     install_requires=[
36         f"jax{_jax_version_constraints}",
37         f"jaxlib{_jaxlib_version_constraints}",
38         "multipledispatch",
39         "numpy",
40         "tqdm",
41     ],
42     extras_require={
43         "doc": [
44             "ipython",  # sphinx needs this to render codes
45             "nbsphinx>=0.8.5",
46             "readthedocs-sphinx-search==0.1.0",
47             "sphinx",
48             "sphinx_rtd_theme",
49             "sphinx-gallery",
50         ],
51         "test": [
52             "black[jupyter]>=21.8b0",
53             "flake8",
54             "isort>=5.0",
55             "pytest>=4.1",
56             "pyro-api>=0.1.1",
57             "scipy>=1.6,<1.7",
58         ],
59         "dev": [
60             "dm-haiku",
61             "flax",
62             "funsor>=0.4.1",
63             "graphviz",
64             "jaxns==0.0.7",
65             "optax>=0.0.6",
66             "tensorflow_probability>=0.15.0",
67         ],
68         "examples": [
69             "arviz",
70             "jupyter",
71             "matplotlib",
72             "pandas",
73             "seaborn",
74             "scikit-learn",
75             "wordcloud",
76         ],
77         "cpu": f"jax[cpu]{_jax_version_constraints}",
78         # TPU and CUDA installations, currently require to add package repository URL, i.e.,
79         # pip install numpyro[cuda] -f https://storage.googleapis.com/jax-releases/jax_releases.html
80         "tpu": f"jax[tpu]{_jax_version_constraints}",
81         "cuda": f"jax[cuda]{_jax_version_constraints}",
82     },
83     long_description=long_description,
84     long_description_content_type="text/markdown",
85     keywords="probabilistic machine learning bayesian statistics",
86     license="Apache License 2.0",
87     classifiers=[
88         "Intended Audience :: Developers",
89         "Intended Audience :: Education",
90         "Intended Audience :: Science/Research",
91         "License :: OSI Approved :: Apache Software License",
92         "Operating System :: POSIX :: Linux",
93         "Operating System :: MacOS :: MacOS X",
94         "Programming Language :: Python :: 3.7",
95         "Programming Language :: Python :: 3.8",
96         "Programming Language :: Python :: 3.9",
97         "Programming Language :: Python :: 3.10",
98     ],
99 )
100
```

Path: `numpyro/ops/provenance.py`
Content:
```
1 # Copyright Contributors to the Pyro project.
2 # SPDX-License-Identifier: Apache-2.0
3 
4 import jax
5 from jax.interpreters import partial_eval
6 from jax.linear_util import wrap_init
7 
8 
9 class _ProvenanceJaxprTrace(partial_eval.DynamicJaxprTrace):
10     """A JAX class to control the behavior of primitives on tracers."""
11 
12     def process_primitive(self, primitive, tracers, params):
13         # remove "_provenance" dimension in arguments before executing the function
14         provenances = [
15             t.aval.named_shape.pop("_provenance", frozenset()) for t in tracers
16         ]
17         out_tracers = super().process_primitive(primitive, tracers, params)
18         # add "_provenance" dimensions to arguments again
19         for t, p in zip(tracers, provenances):
20             if p:
21                 t.aval.named_shape["_provenance"] = p
22 
23         # update outputs' provenance
24         out_provenance = frozenset().union(*provenances)
25         if out_provenance:
26             out_tracers = out_tracers if primitive.multiple_results else [out_tracers]
27             for t in out_tracers:
28                 t.aval.named_shape["_provenance"] = out_provenance
29             out_tracers = out_tracers if primitive.multiple_results else out_tracers[0]
30         return out_tracers
31 
32 
33 class ProvenanceArray:
34     """
35     Provenance tracking implementation in JAX.
36 
37     This class wraps an ndarray to track provenance through JAX ops,
38     where provenance is a user-defined frozenset of objects. The
39     provenance of the output arrays of any op is the union of provenances
40     of input arrays.
41 
42     -   To start tracking provenance in a function, wrap input arrays in
43         :class:`ProvenanceArray` with user-defined initial provenance,
44         then use :func:`eval_provenance` to get the provenance output array.
45     -   To read the provenance of an ndarray use :func:`get_provenance` .
46 
47     Example::
48 
49         >>> a = ProvenanceArray(jnp.zeros(3), frozenset({"a"}))
50         >>> b = ProvenanceArray(jnp.ones(3), frozenset({"b"}))
51         >>> c = jnp.arange(3)
52         >>> f = lambda a, b, c: a + b + c
53         >>> o = eval_provenance(f, a, b, c)
54         >>> assert get_provenance(o) == frozenset({"a", "b"})
55 
56     **References**
57 
58     [1] David Wingate, Noah Goodman, Andreas Stuhlmüller, Jeffrey Siskind (2011)
59         Nonstandard Interpretations of Probabilistic Programs for Efficient Inference
60         http://papers.neurips.cc/paper/4309-nonstandard-interpretations-of-probabilistic-programs-for-efficient-inference.pdf
61 
62     :param data: An initial data to start tracking. The data needs
63         to have attributes `shape` and `dtype`.
64     :param frozenset provenance: An initial provenance set.
65     """
66 
67     def __init__(self, data, provenance=frozenset()):
68         self.shape = data.shape
69         self.dtype = data.dtype
70         self.named_shape = {"_provenance": provenance}
71 
72 
73 def get_provenance(data):
74     """
75     Reads the provenance of a recursive datastructure possibly containing ndarray.
76 
77     :param data: An input data.
78     :returns: A provenance frozenset.
79     :rtype: frozenset
80     """
81     return jax.tree_util.tree_map(
82         lambda a: a.named_shape.get("_provenance", frozenset()), data
83     )
84 
85 
86 def eval_provenance(fun, *args, **kwargs):
87     """
88     Compute the provenance output of ``fun`` using JAX's abstract
89     interpretation machinery. There is no actual array computation performed.
90 
91     :param fun: A callable to track provenance of its (keyword) arguments.
92     :param args: Positional arguments of `fun`.
93     :param kwargs: Keyword arguments of `fun`.
94     :returns: A pytree of :class:`ProvenanceArray`.
95     """
96     # flatten the function and its arguments
97     args_flat, in_tree = jax.tree_util.tree_flatten((args, kwargs))
98     wrapped_fun, out_tree = jax.api_util.flatten_fun(wrap_init(fun), in_tree)
99     fun = wrap_init(wrapped_fun.call_wrapped)
100     avals = jax.util.safe_map(jax.api_util.shaped_abstractify, args_flat)
101 
102     # execute the function and trace provenance
103     with jax.core.new_main(_ProvenanceJaxprTrace, dynamic=True) as main:
104         main.jaxpr_stack = ()
105         out = partial_eval.trace_to_subjaxpr_dynamic(fun, main, avals)[1]
106 
107     # unflatten the output and get its provenance
108     out = [jax.ShapeDtypeStruct(x.shape, x.dtype, x.named_shape) for x in out]
109     out = jax.tree_util.tree_unflatten(out_tree(), out)
110     return jax.tree_util.tree_map(
111         lambda x: ProvenanceArray(x, x.named_shape.get("_provenance", frozenset())),
112         out,
113     )
114
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/numpyro/ops/provenance.py b/numpyro/ops/provenance.py
--- a/numpyro/ops/provenance.py
+++ b/numpyro/ops/provenance.py
@@ -26,6 +26,9 @@
             out_tracers = out_tracers if primitive.multiple_results else [out_tracers]
             for t in out_tracers:
                 t.aval.named_shape["_provenance"] = out_provenance
+                # Also update provenance of the cached tracer -> aval dict.
+                aval_cache = self.frame.tracer_to_var[id(t)].aval
+                aval_cache.named_shape["_provenance"] = out_provenance
             out_tracers = out_tracers if primitive.multiple_results else out_tracers[0]
         return out_tracers
 
diff --git a/setup.py b/setup.py
--- a/setup.py
+++ b/setup.py
@@ -9,7 +9,7 @@
 from setuptools import find_packages, setup
 
 PROJECT_PATH = os.path.dirname(os.path.abspath(__file__))
-_jax_version_constraints = ">=0.2.13,<0.2.28"
+_jax_version_constraints = ">=0.2.13"
 _jaxlib_version_constraints = ">=0.1.65"
 
 # Find version

{"golden_diff": "diff --git a/numpyro/ops/provenance.py b/numpyro/ops/provenance.py\n--- a/numpyro/ops/provenance.py\n+++ b/numpyro/ops/provenance.py\n@@ -26,6 +26,9 @@\n             out_tracers = out_tracers if primitive.multiple_results else [out_tracers]\n             for t in out_tracers:\n                 t.aval.named_shape[\"_provenance\"] = out_provenance\n+                # Also update provenance of the cached tracer -> aval dict.\n+                aval_cache = self.frame.tracer_to_var[id(t)].aval\n+                aval_cache.named_shape[\"_provenance\"] = out_provenance\n             out_tracers = out_tracers if primitive.multiple_results else out_tracers[0]\n         return out_tracers\n \ndiff --git a/setup.py b/setup.py\n--- a/setup.py\n+++ b/setup.py\n@@ -9,7 +9,7 @@\n from setuptools import find_packages, setup\n \n PROJECT_PATH = os.path.dirname(os.path.abspath(__file__))\n-_jax_version_constraints = \">=0.2.13,<0.2.28\"\n+_jax_version_constraints = \">=0.2.13\"\n _jaxlib_version_constraints = \">=0.1.65\"\n \n # Find version\n", "issue": "`numpyro.render_model` behaviour changed in 0.9.0\nI noticed that the behaviour of `numpyro.render_model` appears to have changed since the latest release, specifically arrows signifying conditional dependence are not being rendered properly.\r\n\r\nI haven't had time to dig into why this is yet, but here is a simple recreation\r\n\r\n```python\r\nimport numpyro\r\nimport numpyro.distributions as dist\r\n\r\n\r\ndef model():\r\n    theta = numpyro.sample(\"theta\", dist.Uniform())\r\n    \r\n    with numpyro.plate(\"N\", size=10):\r\n        numpyro.sample(\"y\", dist.Bernoulli(theta))\r\n        \r\n        \r\nif __name__ == \"__main__\":\r\n    g = numpyro.render_model(model)\r\n    g.render(f\"model-{numpyro.__version__}\", format=\"png\")\r\n```\r\n\r\nRunning this with 0.8.0 and 0.9.0 I get (respectively)\r\n\r\n![model-0 8 0](https://user-images.githubusercontent.com/15220906/152194532-07f70d39-8197-41c5-9e14-beefc33b807e.png)\r\n\r\n![model-0 9 0](https://user-images.githubusercontent.com/15220906/152194539-98a41b9f-a673-4937-b864-1fb9e6c3d420.png)\r\n\r\n\n", "before_files": [{"content": "# Copyright Contributors to the Pyro project.\n# SPDX-License-Identifier: Apache-2.0\n\nfrom __future__ import absolute_import, division, print_function\n\nimport os\nimport sys\n\nfrom setuptools import find_packages, setup\n\nPROJECT_PATH = os.path.dirname(os.path.abspath(__file__))\n_jax_version_constraints = \">=0.2.13,<0.2.28\"\n_jaxlib_version_constraints = \">=0.1.65\"\n\n# Find version\nfor line in open(os.path.join(PROJECT_PATH, \"numpyro\", \"version.py\")):\n    if line.startswith(\"__version__ = \"):\n        version = line.strip().split()[2][1:-1]\n\n# READ README.md for long description on PyPi.\ntry:\n    long_description = open(\"README.md\", encoding=\"utf-8\").read()\nexcept Exception as e:\n    sys.stderr.write(\"Failed to read README.md:\\n  {}\\n\".format(e))\n    sys.stderr.flush()\n    long_description = \"\"\n\nsetup(\n    name=\"numpyro\",\n    version=version,\n    description=\"Pyro PPL on NumPy\",\n    packages=find_packages(include=[\"numpyro\", \"numpyro.*\"]),\n    url=\"https://github.com/pyro-ppl/numpyro\",\n    author=\"Uber AI Labs\",\n    install_requires=[\n        f\"jax{_jax_version_constraints}\",\n        f\"jaxlib{_jaxlib_version_constraints}\",\n        \"multipledispatch\",\n        \"numpy\",\n        \"tqdm\",\n    ],\n    extras_require={\n        \"doc\": [\n            \"ipython\",  # sphinx needs this to render codes\n            \"nbsphinx>=0.8.5\",\n            \"readthedocs-sphinx-search==0.1.0\",\n            \"sphinx\",\n            \"sphinx_rtd_theme\",\n            \"sphinx-gallery\",\n        ],\n        \"test\": [\n            \"black[jupyter]>=21.8b0\",\n            \"flake8\",\n            \"isort>=5.0\",\n            \"pytest>=4.1\",\n            \"pyro-api>=0.1.1\",\n            \"scipy>=1.6,<1.7\",\n        ],\n        \"dev\": [\n            \"dm-haiku\",\n            \"flax\",\n            \"funsor>=0.4.1\",\n            \"graphviz\",\n            \"jaxns==0.0.7\",\n            \"optax>=0.0.6\",\n            \"tensorflow_probability>=0.15.0\",\n        ],\n        \"examples\": [\n            \"arviz\",\n            \"jupyter\",\n            \"matplotlib\",\n            \"pandas\",\n            \"seaborn\",\n            \"scikit-learn\",\n            \"wordcloud\",\n        ],\n        \"cpu\": f\"jax[cpu]{_jax_version_constraints}\",\n        # TPU and CUDA installations, currently require to add package repository URL, i.e.,\n        # pip install numpyro[cuda] -f https://storage.googleapis.com/jax-releases/jax_releases.html\n        \"tpu\": f\"jax[tpu]{_jax_version_constraints}\",\n        \"cuda\": f\"jax[cuda]{_jax_version_constraints}\",\n    },\n    long_description=long_description,\n    long_description_content_type=\"text/markdown\",\n    keywords=\"probabilistic machine learning bayesian statistics\",\n    license=\"Apache License 2.0\",\n    classifiers=[\n        \"Intended Audience :: Developers\",\n        \"Intended Audience :: Education\",\n        \"Intended Audience :: Science/Research\",\n        \"License :: OSI Approved :: Apache Software License\",\n        \"Operating System :: POSIX :: Linux\",\n        \"Operating System :: MacOS :: MacOS X\",\n        \"Programming Language :: Python :: 3.7\",\n        \"Programming Language :: Python :: 3.8\",\n        \"Programming Language :: Python :: 3.9\",\n        \"Programming Language :: Python :: 3.10\",\n    ],\n)\n", "path": "setup.py"}, {"content": "# Copyright Contributors to the Pyro project.\n# SPDX-License-Identifier: Apache-2.0\n\nimport jax\nfrom jax.interpreters import partial_eval\nfrom jax.linear_util import wrap_init\n\n\nclass _ProvenanceJaxprTrace(partial_eval.DynamicJaxprTrace):\n    \"\"\"A JAX class to control the behavior of primitives on tracers.\"\"\"\n\n    def process_primitive(self, primitive, tracers, params):\n        # remove \"_provenance\" dimension in arguments before executing the function\n        provenances = [\n            t.aval.named_shape.pop(\"_provenance\", frozenset()) for t in tracers\n        ]\n        out_tracers = super().process_primitive(primitive, tracers, params)\n        # add \"_provenance\" dimensions to arguments again\n        for t, p in zip(tracers, provenances):\n            if p:\n                t.aval.named_shape[\"_provenance\"] = p\n\n        # update outputs' provenance\n        out_provenance = frozenset().union(*provenances)\n        if out_provenance:\n            out_tracers = out_tracers if primitive.multiple_results else [out_tracers]\n            for t in out_tracers:\n                t.aval.named_shape[\"_provenance\"] = out_provenance\n            out_tracers = out_tracers if primitive.multiple_results else out_tracers[0]\n        return out_tracers\n\n\nclass ProvenanceArray:\n    \"\"\"\n    Provenance tracking implementation in JAX.\n\n    This class wraps an ndarray to track provenance through JAX ops,\n    where provenance is a user-defined frozenset of objects. The\n    provenance of the output arrays of any op is the union of provenances\n    of input arrays.\n\n    -   To start tracking provenance in a function, wrap input arrays in\n        :class:`ProvenanceArray` with user-defined initial provenance,\n        then use :func:`eval_provenance` to get the provenance output array.\n    -   To read the provenance of an ndarray use :func:`get_provenance` .\n\n    Example::\n\n        >>> a = ProvenanceArray(jnp.zeros(3), frozenset({\"a\"}))\n        >>> b = ProvenanceArray(jnp.ones(3), frozenset({\"b\"}))\n        >>> c = jnp.arange(3)\n        >>> f = lambda a, b, c: a + b + c\n        >>> o = eval_provenance(f, a, b, c)\n        >>> assert get_provenance(o) == frozenset({\"a\", \"b\"})\n\n    **References**\n\n    [1] David Wingate, Noah Goodman, Andreas Stuhlm\u00fcller, Jeffrey Siskind (2011)\n        Nonstandard Interpretations of Probabilistic Programs for Efficient Inference\n        http://papers.neurips.cc/paper/4309-nonstandard-interpretations-of-probabilistic-programs-for-efficient-inference.pdf\n\n    :param data: An initial data to start tracking. The data needs\n        to have attributes `shape` and `dtype`.\n    :param frozenset provenance: An initial provenance set.\n    \"\"\"\n\n    def __init__(self, data, provenance=frozenset()):\n        self.shape = data.shape\n        self.dtype = data.dtype\n        self.named_shape = {\"_provenance\": provenance}\n\n\ndef get_provenance(data):\n    \"\"\"\n    Reads the provenance of a recursive datastructure possibly containing ndarray.\n\n    :param data: An input data.\n    :returns: A provenance frozenset.\n    :rtype: frozenset\n    \"\"\"\n    return jax.tree_util.tree_map(\n        lambda a: a.named_shape.get(\"_provenance\", frozenset()), data\n    )\n\n\ndef eval_provenance(fun, *args, **kwargs):\n    \"\"\"\n    Compute the provenance output of ``fun`` using JAX's abstract\n    interpretation machinery. There is no actual array computation performed.\n\n    :param fun: A callable to track provenance of its (keyword) arguments.\n    :param args: Positional arguments of `fun`.\n    :param kwargs: Keyword arguments of `fun`.\n    :returns: A pytree of :class:`ProvenanceArray`.\n    \"\"\"\n    # flatten the function and its arguments\n    args_flat, in_tree = jax.tree_util.tree_flatten((args, kwargs))\n    wrapped_fun, out_tree = jax.api_util.flatten_fun(wrap_init(fun), in_tree)\n    fun = wrap_init(wrapped_fun.call_wrapped)\n    avals = jax.util.safe_map(jax.api_util.shaped_abstractify, args_flat)\n\n    # execute the function and trace provenance\n    with jax.core.new_main(_ProvenanceJaxprTrace, dynamic=True) as main:\n        main.jaxpr_stack = ()\n        out = partial_eval.trace_to_subjaxpr_dynamic(fun, main, avals)[1]\n\n    # unflatten the output and get its provenance\n    out = [jax.ShapeDtypeStruct(x.shape, x.dtype, x.named_shape) for x in out]\n    out = jax.tree_util.tree_unflatten(out_tree(), out)\n    return jax.tree_util.tree_map(\n        lambda x: ProvenanceArray(x, x.named_shape.get(\"_provenance\", frozenset())),\n        out,\n    )\n", "path": "numpyro/ops/provenance.py"}], "after_files": [{"content": "# Copyright Contributors to the Pyro project.\n# SPDX-License-Identifier: Apache-2.0\n\nfrom __future__ import absolute_import, division, print_function\n\nimport os\nimport sys\n\nfrom setuptools import find_packages, setup\n\nPROJECT_PATH = os.path.dirname(os.path.abspath(__file__))\n_jax_version_constraints = \">=0.2.13\"\n_jaxlib_version_constraints = \">=0.1.65\"\n\n# Find version\nfor line in open(os.path.join(PROJECT_PATH, \"numpyro\", \"version.py\")):\n    if line.startswith(\"__version__ = \"):\n        version = line.strip().split()[2][1:-1]\n\n# READ README.md for long description on PyPi.\ntry:\n    long_description = open(\"README.md\", encoding=\"utf-8\").read()\nexcept Exception as e:\n    sys.stderr.write(\"Failed to read README.md:\\n  {}\\n\".format(e))\n    sys.stderr.flush()\n    long_description = \"\"\n\nsetup(\n    name=\"numpyro\",\n    version=version,\n    description=\"Pyro PPL on NumPy\",\n    packages=find_packages(include=[\"numpyro\", \"numpyro.*\"]),\n    url=\"https://github.com/pyro-ppl/numpyro\",\n    author=\"Uber AI Labs\",\n    install_requires=[\n        f\"jax{_jax_version_constraints}\",\n        f\"jaxlib{_jaxlib_version_constraints}\",\n        \"multipledispatch\",\n        \"numpy\",\n        \"tqdm\",\n    ],\n    extras_require={\n        \"doc\": [\n            \"ipython\",  # sphinx needs this to render codes\n            \"nbsphinx>=0.8.5\",\n            \"readthedocs-sphinx-search==0.1.0\",\n            \"sphinx\",\n            \"sphinx_rtd_theme\",\n            \"sphinx-gallery\",\n        ],\n        \"test\": [\n            \"black[jupyter]>=21.8b0\",\n            \"flake8\",\n            \"isort>=5.0\",\n            \"pytest>=4.1\",\n            \"pyro-api>=0.1.1\",\n            \"scipy>=1.6,<1.7\",\n        ],\n        \"dev\": [\n            \"dm-haiku\",\n            \"flax\",\n            \"funsor>=0.4.1\",\n            \"graphviz\",\n            \"jaxns==0.0.7\",\n            \"optax>=0.0.6\",\n            \"tensorflow_probability>=0.15.0\",\n        ],\n        \"examples\": [\n            \"arviz\",\n            \"jupyter\",\n            \"matplotlib\",\n            \"pandas\",\n            \"seaborn\",\n            \"scikit-learn\",\n            \"wordcloud\",\n        ],\n        \"cpu\": f\"jax[cpu]{_jax_version_constraints}\",\n        # TPU and CUDA installations, currently require to add package repository URL, i.e.,\n        # pip install numpyro[cuda] -f https://storage.googleapis.com/jax-releases/jax_releases.html\n        \"tpu\": f\"jax[tpu]{_jax_version_constraints}\",\n        \"cuda\": f\"jax[cuda]{_jax_version_constraints}\",\n    },\n    long_description=long_description,\n    long_description_content_type=\"text/markdown\",\n    keywords=\"probabilistic machine learning bayesian statistics\",\n    license=\"Apache License 2.0\",\n    classifiers=[\n        \"Intended Audience :: Developers\",\n        \"Intended Audience :: Education\",\n        \"Intended Audience :: Science/Research\",\n        \"License :: OSI Approved :: Apache Software License\",\n        \"Operating System :: POSIX :: Linux\",\n        \"Operating System :: MacOS :: MacOS X\",\n        \"Programming Language :: Python :: 3.7\",\n        \"Programming Language :: Python :: 3.8\",\n        \"Programming Language :: Python :: 3.9\",\n        \"Programming Language :: Python :: 3.10\",\n    ],\n)\n", "path": "setup.py"}, {"content": "# Copyright Contributors to the Pyro project.\n# SPDX-License-Identifier: Apache-2.0\n\nimport jax\nfrom jax.interpreters import partial_eval\nfrom jax.linear_util import wrap_init\n\n\nclass _ProvenanceJaxprTrace(partial_eval.DynamicJaxprTrace):\n    \"\"\"A JAX class to control the behavior of primitives on tracers.\"\"\"\n\n    def process_primitive(self, primitive, tracers, params):\n        # remove \"_provenance\" dimension in arguments before executing the function\n        provenances = [\n            t.aval.named_shape.pop(\"_provenance\", frozenset()) for t in tracers\n        ]\n        out_tracers = super().process_primitive(primitive, tracers, params)\n        # add \"_provenance\" dimensions to arguments again\n        for t, p in zip(tracers, provenances):\n            if p:\n                t.aval.named_shape[\"_provenance\"] = p\n\n        # update outputs' provenance\n        out_provenance = frozenset().union(*provenances)\n        if out_provenance:\n            out_tracers = out_tracers if primitive.multiple_results else [out_tracers]\n            for t in out_tracers:\n                t.aval.named_shape[\"_provenance\"] = out_provenance\n                # Also update provenance of the cached tracer -> aval dict.\n                aval_cache = self.frame.tracer_to_var[id(t)].aval\n                aval_cache.named_shape[\"_provenance\"] = out_provenance\n            out_tracers = out_tracers if primitive.multiple_results else out_tracers[0]\n        return out_tracers\n\n\nclass ProvenanceArray:\n    \"\"\"\n    Provenance tracking implementation in JAX.\n\n    This class wraps an ndarray to track provenance through JAX ops,\n    where provenance is a user-defined frozenset of objects. The\n    provenance of the output arrays of any op is the union of provenances\n    of input arrays.\n\n    -   To start tracking provenance in a function, wrap input arrays in\n        :class:`ProvenanceArray` with user-defined initial provenance,\n        then use :func:`eval_provenance` to get the provenance output array.\n    -   To read the provenance of an ndarray use :func:`get_provenance` .\n\n    Example::\n\n        >>> a = ProvenanceArray(jnp.zeros(3), frozenset({\"a\"}))\n        >>> b = ProvenanceArray(jnp.ones(3), frozenset({\"b\"}))\n        >>> c = jnp.arange(3)\n        >>> f = lambda a, b, c: a + b + c\n        >>> o = eval_provenance(f, a, b, c)\n        >>> assert get_provenance(o) == frozenset({\"a\", \"b\"})\n\n    **References**\n\n    [1] David Wingate, Noah Goodman, Andreas Stuhlm\u00fcller, Jeffrey Siskind (2011)\n        Nonstandard Interpretations of Probabilistic Programs for Efficient Inference\n        http://papers.neurips.cc/paper/4309-nonstandard-interpretations-of-probabilistic-programs-for-efficient-inference.pdf\n\n    :param data: An initial data to start tracking. The data needs\n        to have attributes `shape` and `dtype`.\n    :param frozenset provenance: An initial provenance set.\n    \"\"\"\n\n    def __init__(self, data, provenance=frozenset()):\n        self.shape = data.shape\n        self.dtype = data.dtype\n        self.named_shape = {\"_provenance\": provenance}\n\n\ndef get_provenance(data):\n    \"\"\"\n    Reads the provenance of a recursive datastructure possibly containing ndarray.\n\n    :param data: An input data.\n    :returns: A provenance frozenset.\n    :rtype: frozenset\n    \"\"\"\n    return jax.tree_util.tree_map(\n        lambda a: a.named_shape.get(\"_provenance\", frozenset()), data\n    )\n\n\ndef eval_provenance(fun, *args, **kwargs):\n    \"\"\"\n    Compute the provenance output of ``fun`` using JAX's abstract\n    interpretation machinery. There is no actual array computation performed.\n\n    :param fun: A callable to track provenance of its (keyword) arguments.\n    :param args: Positional arguments of `fun`.\n    :param kwargs: Keyword arguments of `fun`.\n    :returns: A pytree of :class:`ProvenanceArray`.\n    \"\"\"\n    # flatten the function and its arguments\n    args_flat, in_tree = jax.tree_util.tree_flatten((args, kwargs))\n    wrapped_fun, out_tree = jax.api_util.flatten_fun(wrap_init(fun), in_tree)\n    fun = wrap_init(wrapped_fun.call_wrapped)\n    avals = jax.util.safe_map(jax.api_util.shaped_abstractify, args_flat)\n\n    # execute the function and trace provenance\n    with jax.core.new_main(_ProvenanceJaxprTrace, dynamic=True) as main:\n        main.jaxpr_stack = ()\n        out = partial_eval.trace_to_subjaxpr_dynamic(fun, main, avals)[1]\n\n    # unflatten the output and get its provenance\n    out = [jax.ShapeDtypeStruct(x.shape, x.dtype, x.named_shape) for x in out]\n    out = jax.tree_util.tree_unflatten(out_tree(), out)\n    return jax.tree_util.tree_map(\n        lambda x: ProvenanceArray(x, x.named_shape.get(\"_provenance\", frozenset())),\n        out,\n    )\n", "path": "numpyro/ops/provenance.py"}]}