code: def _extract_provenance_pytree(x):
      flat_args, spec = tree_flatten(x)
      xs = []
      provenance = frozenset()
      for x, p in map(extract_provenance, flat_args):
          xs.append(x)
          provenance |= p
      return tree_unflatten(xs, spec), provenance
  
file_path: pyro/pyro/ops/provenance.py
imports: from functools import partial, singledispatch
  from typing import Tuple
  import torch
  from torch.utils._pytree import tree_flatten, tree_map, tree_unflatten
label_desc: Use image node_id 8 for calling a global function with example usage: _extract_provenance_pytree(x) and returns: tree_unflatten,  provenance