Repository: pyro-ppl/numpyro

Issue:
`NameError: unbound axis name: _provenance` with the latest jax version
I'm getting the error message below when running `pytest test/contrib/test_enum_elbo.py` with the latest version of jax (0.4.4):

```
self = Traced<ShapedArray(int32[1;_provenance:frozenset({'w'})])>with<DynamicJaxprTrace(level=4/0)>, other = 0

    def deferring_binary_op(self, other):
      if hasattr(other, '__jax_array__'):
        other = other.__jax_array__()
      args = (other, self) if swap else (self, other)
      if isinstance(other, _accepted_binop_types):
>       return binary_op(*args)
E       NameError: unbound axis name: _provenance. The following axis names (e.g. defined by pmap) are available to collective operations: []

/home/yordabay/anaconda3/envs/numpyro/lib/python3.8/site-packages/jax/_src/numpy/lax_numpy.py:5071: NameError
```


Discussion:

Resolving patch:
diff --git a/numpyro/infer/elbo.py b/numpyro/infer/elbo.py
--- a/numpyro/infer/elbo.py
+++ b/numpyro/infer/elbo.py
@@ -14,14 +14,14 @@
 from numpyro.distributions import ExpandedDistribution, MaskedDistribution
 from numpyro.distributions.kl import kl_divergence
 from numpyro.distributions.util import scale_and_mask
-from numpyro.handlers import Messenger, replay, seed, substitute, trace
+from numpyro.handlers import replay, seed, substitute, trace
 from numpyro.infer.util import (
     _without_rsample_stop_gradient,
     get_importance_trace,
     is_identically_one,
     log_density,
 )
-from numpyro.ops.provenance import eval_provenance, get_provenance
+from numpyro.ops.provenance import eval_provenance
 from numpyro.util import _validate_model, check_model_guide_match, find_stack_level
 
 
@@ -535,63 +535,6 @@ def _compute_downstream_costs(model_trace, guide_trace, non_reparam_nodes):
     return downstream_costs, downstream_guide_cost_nodes
 
 
-class track_nonreparam(Messenger):
-    """
-    Track non-reparameterizable sample sites. Intended to be used with ``eval_provenance``.
-
-    **References:**
-
-    1. *Nonstandard Interpretations of Probabilistic Programs for Efficient Inference*,
-        David Wingate, Noah Goodman, Andreas Stuhlmüller, Jeffrey Siskind
-
-    **Example:**
-
-    .. doctest::
-
-       >>> import jax.numpy as jnp
-       >>> import numpyro
-       >>> import numpyro.distributions as dist
-       >>> from numpyro.infer.elbo import track_nonreparam
-       >>> from numpyro.ops.provenance import eval_provenance, get_provenance
-       >>> from numpyro.handlers import seed, trace
-
-       >>> def model():
-       ...     probs_a = jnp.array([0.3, 0.7])
-       ...     probs_b = jnp.array([[0.1, 0.9], [0.8, 0.2]])
-       ...     probs_c = jnp.array([[0.5, 0.5], [0.6, 0.4]])
-       ...     a = numpyro.sample("a", dist.Categorical(probs_a))
-       ...     b = numpyro.sample("b", dist.Categorical(probs_b[a]))
-       ...     numpyro.sample("c", dist.Categorical(probs_c[b]), obs=jnp.array(0))
-
-       >>> def get_log_probs():
-       ...     seeded_model = seed(model, rng_seed=0)
-       ...     model_tr = trace(seeded_model).get_trace()
-       ...     return {
-       ...         name: site["fn"].log_prob(site["value"])
-       ...         for name, site in model_tr.items()
-       ...         if site["type"] == "sample"
-       ...     }
-
-       >>> model_deps = get_provenance(eval_provenance(track_nonreparam(get_log_probs)))
-       >>> print(model_deps)  # doctest: +SKIP
-       {'a': frozenset({'a'}), 'b': frozenset({'a', 'b'}), 'c': frozenset({'a', 'b'})}
-    """
-
-    def postprocess_message(self, msg):
-        if (
-            msg["type"] == "sample"
-            and (not msg["is_observed"])
-            and (not msg["fn"].has_rsample)
-        ):
-            new_provenance = frozenset({msg["name"]})
-            old_provenance = msg["value"].aval.named_shape.get(
-                "_provenance", frozenset()
-            )
-            msg["value"].aval.named_shape["_provenance"] = (
-                old_provenance | new_provenance
-            )
-
-
 def get_importance_log_probs(model, guide, args, kwargs, params):
     """
     Returns log probabilities at each site for the guide and the model that is run against it.
@@ -610,6 +553,43 @@ def get_importance_log_probs(model, guide, args, kwargs, params):
     return model_log_probs, guide_log_probs
 
 
+def _substitute_nonreparam(data, msg):
+    if msg["name"] in data and not msg["fn"].has_rsample:
+        value = msg["fn"](*msg["args"], **msg["kwargs"])
+        value = 0 * value + data[msg["name"]]
+        return value
+
+
+def _get_latents(model, guide, args, kwargs, params):
+    model = seed(substitute(model, data=params), rng_seed=0)
+    guide = seed(substitute(guide, data=params), rng_seed=0)
+    guide_tr = trace(guide).get_trace(*args, **kwargs)
+    model_tr = trace(replay(model, guide_tr)).get_trace(*args, **kwargs)
+    model_tr.update(guide_tr)
+    return {
+        name: site["value"]
+        for name, site in model_tr.items()
+        if site["type"] == "sample" and not site.get("is_observed", False)
+    }
+
+
+def get_nonreparam_deps(model, guide, args, kwargs, param_map, latents=None):
+    """Find dependencies on non-reparameterizable sample sites for each cost term in the model and the guide."""
+    if latents is None:
+        latents = eval_shape(
+            partial(_get_latents, model, guide, args, kwargs, param_map)
+        )
+
+    def fn(**latents):
+        subs_fn = partial(_substitute_nonreparam, latents)
+        subs_model = substitute(seed(model, rng_seed=0), substitute_fn=subs_fn)
+        subs_guide = substitute(seed(guide, rng_seed=0), substitute_fn=subs_fn)
+        return get_importance_log_probs(subs_model, subs_guide, args, kwargs, param_map)
+
+    model_deps, guide_deps = eval_provenance(fn, **latents)
+    return model_deps, guide_deps
+
+
 class TraceGraph_ELBO(ELBO):
     """
     A TraceGraph implementation of ELBO-based SVI. The gradient estimator
@@ -661,8 +641,12 @@ def single_particle_elbo(rng_key):
             check_model_guide_match(model_trace, guide_trace)
             _validate_model(model_trace, plate_warning="strict")
 
+            latents = {}
+            for name, site in guide_trace.items():
+                if site["type"] == "sample" and not site.get("is_observed", False):
+                    latents[name] = site["value"]
             model_deps, guide_deps = get_nonreparam_deps(
-                seeded_model, seeded_guide, args, kwargs, param_map
+                model, guide, args, kwargs, param_map, latents=latents
             )
 
             elbo = 0.0
@@ -837,23 +821,6 @@ def _partition(model_sum_deps, sum_vars):
     return components
 
 
-def get_nonreparam_deps(model, guide, args, kwargs, param_map):
-    """Find dependencies on non-reparameterizable sample sites for each cost term in the model and the guide."""
-    model_deps, guide_deps = get_provenance(
-        eval_provenance(
-            partial(
-                track_nonreparam(get_importance_log_probs),
-                model,
-                guide,
-                args,
-                kwargs,
-                param_map,
-            )
-        )
-    )
-    return model_deps, guide_deps
-
-
 def guess_max_plate_nesting(model, guide, args, kwargs, param_map):
     """Guess maximum plate nesting by performing jax shape inference."""
     model_shapes, guide_shapes = eval_shape(
@@ -917,15 +884,15 @@ def single_particle_elbo(rng_key):
             if self.max_plate_nesting == float("inf"):
                 seeded_model = seed(model, model_seed)
                 seeded_guide = seed(guide, guide_seed)
+                # XXX: We can extract abstract latents here such that they
+                # can be reused in get_nonreparam_deps below.
                 self.max_plate_nesting = guess_max_plate_nesting(
                     seeded_model, seeded_guide, args, kwargs, param_map
                 )
 
-            seeded_model = seed(model, model_seed)
-            seeded_guide = seed(guide, guide_seed)
             # get dependencies on nonreparametrizable variables
             model_deps, guide_deps = get_nonreparam_deps(
-                seeded_model, seeded_guide, args, kwargs, param_map
+                model, guide, args, kwargs, param_map
             )
             # get descendants of variables in the guide
             guide_desc = defaultdict(frozenset)
diff --git a/numpyro/infer/inspect.py b/numpyro/infer/inspect.py
--- a/numpyro/infer/inspect.py
+++ b/numpyro/infer/inspect.py
@@ -11,7 +11,7 @@
 from numpyro import handlers
 import numpyro.distributions as dist
 from numpyro.infer.initialization import init_to_sample
-from numpyro.ops.provenance import ProvenanceArray, eval_provenance, get_provenance
+from numpyro.ops.provenance import eval_provenance
 from numpyro.ops.pytree import PytreeTrace
 
 
@@ -55,7 +55,7 @@ def get_trace():
     return jax.eval_shape(get_trace).trace
 
 
-def _get_log_probs(model, model_args, model_kwargs, sample):
+def _get_log_probs(model, model_args, model_kwargs, **sample):
     # Note: We use seed 0 for parameter initialization.
     with handlers.trace() as tr, handlers.seed(rng_seed=0), handlers.substitute(
         data=sample
@@ -199,14 +199,12 @@ def model_3():
 
     # Find direct prior dependencies among latent and observed sites.
     samples = {
-        name: ProvenanceArray(site["value"], frozenset({name}))
+        name: site["value"]
         for name, site in trace.items()
         if site["type"] == "sample" and not site["is_observed"]
     }
-    sample_deps = get_provenance(
-        eval_provenance(
-            partial(_get_log_probs, model, model_args, model_kwargs), samples
-        )
+    sample_deps = eval_provenance(
+        partial(_get_log_probs, model, model_args, model_kwargs), **samples
     )
     prior_dependencies = {n: {n: set()} for n in plates}  # no deps yet
     for i, downstream in enumerate(sample_sites):
@@ -360,7 +358,7 @@ def _resolve_plate_samples(plate_samples):
         k: [name for name in trace if name in v] for k, v in plate_samples.items()
     }
 
-    def get_log_probs(sample):
+    def get_log_probs(**sample):
         class substitute_deterministic(handlers.substitute):
             def process_message(self, msg):
                 if msg["type"] == "deterministic":
@@ -384,24 +382,17 @@ def process_message(self, msg):
         return provenance_arrays
 
     samples = {
-        name: ProvenanceArray(site["value"], frozenset({name}))
+        name: site["value"]
         for name, site in trace.items()
         if (site["type"] == "sample" and not site["is_observed"])
         or site["type"] == "deterministic"
     }
 
     params = {
-        name: jax.tree_util.tree_map(
-            lambda x: ProvenanceArray(x, frozenset({name})), site["value"]
-        )
-        for name, site in trace.items()
-        if site["type"] == "param"
+        name: site["value"] for name, site in trace.items() if site["type"] == "param"
     }
 
-    sample_and_params = {**samples, **params}
-    sample_params_deps = get_provenance(
-        eval_provenance(get_log_probs, sample_and_params)
-    )
+    sample_params_deps = eval_provenance(get_log_probs, **samples, **params)
 
     sample_sample = {}
     sample_param = {}
diff --git a/numpyro/ops/provenance.py b/numpyro/ops/provenance.py
--- a/numpyro/ops/provenance.py
+++ b/numpyro/ops/provenance.py
@@ -2,116 +2,120 @@
 # SPDX-License-Identifier: Apache-2.0
 
 import jax
-from jax.interpreters import partial_eval
-from jax.linear_util import wrap_init
+from jax._src.pjit import pjit_p
+from jax.api_util import flatten_fun, shaped_abstractify
+import jax.core as core
+from jax.interpreters.partial_eval import debug_info, trace_to_jaxpr_dynamic
+from jax.interpreters.pxla import xla_pmap_p
+from jax.interpreters.xla import xla_call_p
+import jax.linear_util as lu
 import jax.numpy as jnp
 
 
-class _ProvenanceJaxprTrace(partial_eval.DynamicJaxprTrace):
-    """A JAX class to control the behavior of primitives on tracers."""
-
-    def process_primitive(self, primitive, tracers, params):
-        # remove "_provenance" dimension in arguments before executing the function
-        provenances = [
-            t.aval.named_shape.pop("_provenance", frozenset()) for t in tracers
-        ]
-        out_tracers = super().process_primitive(primitive, tracers, params)
-        # add "_provenance" dimensions to arguments again
-        for t, p in zip(tracers, provenances):
-            if p:
-                t.aval.named_shape["_provenance"] = p
-
-        # update outputs' provenance
-        out_provenance = frozenset().union(*provenances)
-        if out_provenance:
-            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
-
-
-class ProvenanceArray:
+def eval_provenance(fn, **kwargs):
     """
-    Provenance tracking implementation in JAX.
-
-    This class wraps an ndarray to track provenance through JAX ops,
-    where provenance is a user-defined frozenset of objects. The
-    provenance of the output arrays of any op is the union of provenances
-    of input arrays.
-
-    -   To start tracking provenance in a function, wrap input arrays in
-        :class:`ProvenanceArray` with user-defined initial provenance,
-        then use :func:`eval_provenance` to get the provenance output array.
-    -   To read the provenance of an ndarray use :func:`get_provenance` .
+    Compute the provenance output of ``fun`` using JAX's abstract
+    interpretation machinery. There is no actual array computation performed.
 
     Example::
 
-        >>> a = ProvenanceArray(jnp.zeros(3), frozenset({"a"}))
-        >>> b = ProvenanceArray(jnp.ones(3), frozenset({"b"}))
-        >>> c = jnp.arange(3)
-        >>> f = lambda a, b, c: a + b + c
-        >>> o = eval_provenance(f, a, b, c)
-        >>> assert get_provenance(o) == frozenset({"a", "b"})
+        >>> o = eval_provenance(lambda x, y, z: x + y, x=1, y=2, z=3)
+        >>> assert o == frozenset({"x", "y"})
 
     **References**
 
     [1] David Wingate, Noah Goodman, Andreas Stuhlmüller, Jeffrey Siskind (2011)
         Nonstandard Interpretations of Probabilistic Programs for Efficient Inference
         http://papers.neurips.cc/paper/4309-nonstandard-interpretations-of-probabilistic-programs-for-efficient-inference.pdf
+    [2] https://jax.readthedocs.io/en/latest/notebooks/Writing_custom_interpreters_in_Jax.html
 
-    :param data: An initial data to start tracking. The data needs
-        to have attributes `shape` and `dtype`.
-    :param frozenset provenance: An initial provenance set.
+    :param fun: A callable to track provenance of its (keyword) arguments.
+    :param kwargs: Keyword arguments of `fun`.
+    :returns: A pytree of :class:`frozenset` indicating the dependency on the inputs.
     """
+    # Flatten the function and its arguments
+    args, in_tree = jax.tree_util.tree_flatten(((), kwargs))
+    wrapped_fun, out_tree = flatten_fun(lu.wrap_init(fn), in_tree)
+    # Abstract eval to get output pytree
+    avals = core.safe_map(shaped_abstractify, args)
+    # XXX: we split out the process of abstract evaluation and provenance tracking
+    # for simplicity. In principle, they can be merged so that we only need to walk
+    # through the equations once.
+    info = debug_info(fn, in_tree, True, "eval_shape")
+    jaxpr, avals_out, _ = trace_to_jaxpr_dynamic(
+        lu.wrap_init(wrapped_fun.call_wrapped, {}), avals, info
+    )
 
-    def __init__(self, data, provenance=frozenset()):
-        self.shape = jnp.shape(data)
-        self.dtype = jnp.result_type(data)
-        self.named_shape = {"_provenance": provenance}
+    # get provenances of flatten kwargs
+    aval_kwargs = {}
+    for n, v in kwargs.items():
+        aval = jax.ShapeDtypeStruct((), jnp.bool_, {"provenance": frozenset({n})})
+        aval_kwargs[n] = jax.tree_util.tree_map(lambda _: aval, v)
+    aval_args, _ = jax.tree_util.tree_flatten(((), aval_kwargs))
+    provenance_inputs = jax.tree_util.tree_map(
+        lambda x: x.named_shape["provenance"], aval_args
+    )
 
+    provenance_outputs = track_deps_jaxpr(jaxpr, provenance_inputs)
+    out_flat = []
+    for v, p in zip(avals_out, provenance_outputs):
+        val = jax.ShapeDtypeStruct(jnp.shape(v), jnp.result_type(v), {"provenance": p})
+        out_flat.append(val)
+    out = jax.tree_util.tree_unflatten(out_tree(), out_flat)
+    return jax.tree_util.tree_map(lambda x: x.named_shape["provenance"], out)
 
-def get_provenance(data):
-    """
-    Reads the provenance of a recursive data structure possibly containing ndarray.
 
-    :param data: An input data.
-    :returns: A provenance frozenset.
-    :rtype: frozenset
-    """
-    return jax.tree_util.tree_map(
-        lambda a: a.named_shape.get("_provenance", frozenset()), data
-    )
+def track_deps_jaxpr(jaxpr, provenance_inputs):
+    # Mapping from variable -> provenance
+    env = {}
 
+    def read(v):
+        if isinstance(v, core.Literal):
+            return frozenset()
+        return env.get(v, frozenset())
 
-def eval_provenance(fun, *args, **kwargs):
-    """
-    Compute the provenance output of ``fun`` using JAX's abstract
-    interpretation machinery. There is no actual array computation performed.
+    def write(v, p):
+        if isinstance(v, core.Literal):
+            return
+        env[v] = read(v) | p
 
-    :param fun: A callable to track provenance of its (keyword) arguments.
-    :param args: Positional arguments of `fun`.
-    :param kwargs: Keyword arguments of `fun`.
-    :returns: A pytree of :class:`ProvenanceArray`.
-    """
-    # flatten the function and its arguments
-    args_flat, in_tree = jax.tree_util.tree_flatten((args, kwargs))
-    wrapped_fun, out_tree = jax.api_util.flatten_fun(wrap_init(fun), in_tree)
-    fun = wrap_init(wrapped_fun.call_wrapped)
-    avals = jax.util.safe_map(jax.api_util.shaped_abstractify, args_flat)
-
-    # execute the function and trace provenance
-    with jax.core.new_main(_ProvenanceJaxprTrace, dynamic=True) as main:
-        main.jaxpr_stack = ()
-        out = partial_eval.trace_to_subjaxpr_dynamic(fun, main, avals)[1]
-
-    # unflatten the output and get its provenance
-    out = [jax.ShapeDtypeStruct(x.shape, x.dtype, x.named_shape) for x in out]
-    out = jax.tree_util.tree_unflatten(out_tree(), out)
-    return jax.tree_util.tree_map(
-        lambda x: ProvenanceArray(x, x.named_shape.get("_provenance", frozenset())),
-        out,
-    )
+    core.safe_map(write, jaxpr.invars, provenance_inputs)
+    for eqn in jaxpr.eqns:
+        provenance_inputs = core.safe_map(read, eqn.invars)
+        rule = track_deps_rules.get(eqn.primitive, _default_track_deps_rules)
+        provenance_outputs = rule(eqn, provenance_inputs)
+        core.safe_map(write, eqn.outvars, provenance_outputs)
+
+    return core.safe_map(read, jaxpr.outvars)
+
+
+track_deps_rules = {}
+
+
+# XXX: Currently, we use default rule for scan_p, cond_p, while_p, remat_p
+def _default_track_deps_rules(eqn, provenance_inputs):
+    provenance_outputs = frozenset().union(*provenance_inputs)
+    return [provenance_outputs] * len(eqn.outvars)
+
+
+def track_deps_call_rule(eqn, provenance_inputs):
+    return track_deps_jaxpr(eqn.params["call_jaxpr"], provenance_inputs)
+
+
+track_deps_rules[core.call_p] = track_deps_call_rule
+track_deps_rules[xla_call_p] = track_deps_call_rule
+track_deps_rules[xla_pmap_p] = track_deps_call_rule
+
+
+def track_deps_closed_call_rule(eqn, provenance_inputs):
+    return track_deps_jaxpr(eqn.params["call_jaxpr"].jaxpr, provenance_inputs)
+
+
+track_deps_rules[core.closed_call_p] = track_deps_closed_call_rule
+
+
+def track_deps_pjit_rule(eqn, provenance_inputs):
+    return track_deps_jaxpr(eqn.params["jaxpr"].jaxpr, provenance_inputs)
+
+
+track_deps_rules[pjit_p] = track_deps_pjit_rule


Test patch:
diff --git a/test/infer/test_compute_downstream_costs.py b/test/infer/test_compute_downstream_costs.py
--- a/test/infer/test_compute_downstream_costs.py
+++ b/test/infer/test_compute_downstream_costs.py
@@ -19,9 +19,8 @@
     _compute_downstream_costs,
     _get_plate_stacks,
     _identify_dense_edges,
-    track_nonreparam,
+    get_nonreparam_deps,
 )
-from numpyro.ops.provenance import eval_provenance, get_provenance
 
 
 def _brute_force_compute_downstream_costs(
@@ -75,13 +74,11 @@ def _brute_force_compute_downstream_costs(
     return downstream_costs, downstream_guide_cost_nodes
 
 
-def _provenance_compute_downstream_costs(model_trace, guide_trace, get_log_probs):
+def _provenance_compute_downstream_costs(
+    model_trace, guide_trace, model_deps, guide_deps
+):
     # replicate the logic from TraceGraph_ELBO
     # additionally compute downstream_guide_cost_nodes
-    model_deps, guide_deps = get_provenance(
-        eval_provenance(track_nonreparam(get_log_probs))
-    )
-
     downstream_costs = defaultdict(lambda: MultiFrameTensor())
     downstream_guide_cost_nodes = defaultdict(lambda: set())
     for name, site in model_trace.items():
@@ -104,21 +101,6 @@ def _provenance_compute_downstream_costs(model_trace, guide_trace, get_log_probs
     return downstream_costs, downstream_guide_cost_nodes
 
 
-def _get_log_probs(get_traces_fn):
-    model_tr, guide_tr = get_traces_fn()
-    model_log_probs = {
-        name: site["log_prob"]
-        for name, site in model_tr.items()
-        if site["type"] == "sample"
-    }
-    guide_log_probs = {
-        name: site["log_prob"]
-        for name, site in guide_tr.items()
-        if site["type"] == "sample"
-    }
-    return model_log_probs, guide_log_probs
-
-
 def big_model_guide(
     include_obs=True,
     include_single=False,
@@ -201,20 +183,19 @@ def big_model_guide(
 def test_compute_downstream_costs_big_model_guide_pair(
     include_inner_1, include_single, flip_c23, include_triple, include_z1
 ):
-    def _get_traces():
-        seeded_guide = handlers.seed(big_model_guide, rng_seed=0)
-        guide_trace = handlers.trace(seeded_guide).get_trace(
-            include_obs=False,
+    def _get_traces_and_deps():
+        model = partial(
+            big_model_guide,
+            include_obs=True,
             include_inner_1=include_inner_1,
             include_single=include_single,
             flip_c23=flip_c23,
             include_triple=include_triple,
             include_z1=include_z1,
         )
-        model_trace = handlers.trace(
-            handlers.replay(seeded_guide, guide_trace)
-        ).get_trace(
-            include_obs=True,
+        guide = partial(
+            big_model_guide,
+            include_obs=False,
             include_inner_1=include_inner_1,
             include_single=include_single,
             flip_c23=flip_c23,
@@ -222,14 +203,19 @@ def _get_traces():
             include_z1=include_z1,
         )
 
+        seeded_guide = handlers.seed(guide, rng_seed=0)
+        guide_trace = handlers.trace(seeded_guide).get_trace()
+        model_trace = handlers.trace(handlers.replay(model, guide_trace)).get_trace()
+
         for trace in (model_trace, guide_trace):
             for site in trace.values():
                 if site["type"] == "sample":
                     site["log_prob"] = site["fn"].log_prob(site["value"])
 
-        return model_trace, guide_trace
+        model_deps, guide_deps = get_nonreparam_deps(model, guide, (), {}, {})
+        return model_trace, guide_trace, model_deps, guide_deps
 
-    model_trace, guide_trace = _get_traces()
+    model_trace, guide_trace, model_deps, guide_deps = _get_traces_and_deps()
     non_reparam_nodes = set(
         name
         for name, site in guide_trace.items()
@@ -246,7 +232,7 @@ def _get_traces():
     )
 
     dc_provenance, dc_nodes_provenance = _provenance_compute_downstream_costs(
-        model_trace, guide_trace, partial(_get_log_probs, _get_traces)
+        model_trace, guide_trace, model_deps, guide_deps
     )
 
     assert dc_nodes == dc_nodes_brute
@@ -397,23 +383,24 @@ def plate_reuse_model_guide(include_obs=True, dim1=3, dim2=2):
 @pytest.mark.parametrize("dim1", [2, 5])
 @pytest.mark.parametrize("dim2", [3, 4])
 def test_compute_downstream_costs_plate_reuse(dim1, dim2):
-    def _get_traces():
-        seeded_guide = handlers.seed(plate_reuse_model_guide, rng_seed=0)
-        guide_trace = handlers.trace(seeded_guide).get_trace(
-            include_obs=False, dim1=dim1, dim2=dim2
+    def _get_traces_and_deps():
+        model = partial(plate_reuse_model_guide, include_obs=True, dim1=dim1, dim2=dim2)
+        guide = partial(
+            plate_reuse_model_guide, include_obs=False, dim1=dim1, dim2=dim2
         )
-        model_trace = handlers.trace(
-            handlers.replay(seeded_guide, guide_trace)
-        ).get_trace(include_obs=True, dim1=dim1, dim2=dim2)
+        seeded_guide = handlers.seed(guide, rng_seed=0)
+        guide_trace = handlers.trace(seeded_guide).get_trace()
+        model_trace = handlers.trace(handlers.replay(model, guide_trace)).get_trace()
 
         for trace in (model_trace, guide_trace):
             for site in trace.values():
                 if site["type"] == "sample":
                     site["log_prob"] = site["fn"].log_prob(site["value"])
 
-        return model_trace, guide_trace
+        model_deps, guide_deps = get_nonreparam_deps(model, guide, (), {}, {})
+        return model_trace, guide_trace, model_deps, guide_deps
 
-    model_trace, guide_trace = _get_traces()
+    model_trace, guide_trace, model_deps, guide_deps = _get_traces_and_deps()
     non_reparam_nodes = set(
         name
         for name, site in guide_trace.items()
@@ -430,7 +417,7 @@ def _get_traces():
     )
 
     dc_provenance, dc_nodes_provenance = _provenance_compute_downstream_costs(
-        model_trace, guide_trace, partial(_get_log_probs, _get_traces)
+        model_trace, guide_trace, model_deps, guide_deps
     )
 
     assert dc_nodes == dc_nodes_brute
diff --git a/test/ops/test_provenance.py b/test/ops/test_provenance.py
--- a/test/ops/test_provenance.py
+++ b/test/ops/test_provenance.py
@@ -1,21 +1,86 @@
 # Copyright Contributors to the Pyro project.
 # SPDX-License-Identifier: Apache-2.0
 
-import numpy as np
+import inspect
+
 import pytest
 
-from numpyro.ops.provenance import ProvenanceArray, eval_provenance, get_provenance
+import jax
+from jax.api_util import flatten_fun_nokwargs
+import jax.core as core
+import jax.linear_util as lu
+import jax.numpy as jnp
+
+from numpyro.ops.provenance import eval_provenance
 
 
 @pytest.mark.parametrize(
     "f, inputs, expected_output",
     [
-        (lambda x, y: x + y, ({"a"}, {}), {"a"}),
-        (lambda x, y: x + y, ({"a"}, {"b"}), {"a", "b"}),
-        (lambda x, y: x + y, ({"a", "c"}, {"a", "b"}), {"a", "b", "c"}),
-        (lambda x, y, z: x + y, ({"a"}, {"b"}, {"c"}), {"a", "b"}),
+        (lambda a, b: a + 1, ("a", "b"), {"a"}),
+        (lambda a, b: jax.scipy.special.xlogy(a, b), ("a", "b"), {"a", "b"}),
+        (lambda a, b, c: a + b, ("a", "b", "c"), {"a", "b"}),
+        (
+            lambda a, b: {"sum": a + b, "zero": 0},
+            ("a", "b"),
+            {"sum": {"a", "b"}, "zero": set()},
+        ),
     ],
 )
 def test_provenance(f, inputs, expected_output):
-    inputs = [ProvenanceArray(np.array([0]), p) for p in inputs]
-    assert get_provenance(eval_provenance(f, *inputs)) == expected_output
+    inputs = {p: 0 for p in inspect.getfullargspec(f).args}
+    assert eval_provenance(f, **inputs) == expected_output
+
+
+def test_provenance_const():
+    def f(x):
+        with jax.ensure_compile_time_eval():
+            y = jnp.ones(4)
+        return x + y
+
+    jaxpr = jax.make_jaxpr(f)(jnp.zeros((3, 4), jnp.float32))
+    assert len(jaxpr.consts) == 1
+    assert eval_provenance(f, x=3) == {"x"}
+
+
+def test_provenance_fori():
+    def f(x, y, z):
+        del z
+        return jax.lax.fori_loop(0, 5, lambda _, x: x + y, x)
+
+    assert eval_provenance(f, x=3, y=2, z=1) == {"x", "y"}
+
+
+def test_provenance_vmap():
+    def f(x, y):
+        del x
+        return jax.vmap(jnp.sin)(y)
+
+    assert eval_provenance(f, x=3, y=jnp.ones(3)) == {"y"}
+
+
+def test_provenance_pytree_in():
+    def f(x, y):
+        return x["v"] * y, x["u"]
+
+    assert eval_provenance(f, x={"v": 2, "u": 1}, y=1) == ({"x", "y"}, {"x"})
+
+
+def test_provenance_call():
+    def identity(x):
+        args, in_tree = jax.tree_util.tree_flatten((x,))
+        fn, out_tree = flatten_fun_nokwargs(lu.wrap_init(lambda x: x), in_tree)
+        out = core.closed_call_p.bind(fn, *args)
+        return jax.tree_util.tree_unflatten(out_tree(), out)
+
+    assert eval_provenance(identity, x={"v": 2}) == {"v": frozenset({"x"})}
+
+
+def test_provenance_closed_call():
+    def identity(x):
+        args, in_tree = jax.tree_util.tree_flatten((x,))
+        fn, out_tree = flatten_fun_nokwargs(lu.wrap_init(lambda x: x), in_tree)
+        out = core.closed_call_p.bind(fn, *args)
+        return jax.tree_util.tree_unflatten(out_tree(), out)
+
+    assert eval_provenance(identity, x={"v": 2}) == {"v": frozenset({"x"})}