code: def get_model_relations(
      model: Callable,
      model_args: Optional[tuple] = None,
      model_kwargs: Optional[dict] = None,
      include_deterministic: bool = False,
  ):
      """
      Infer relations of RVs and plates from given model and optionally data.
      See https://github.com/pyro-ppl/pyro/issues/949 for more details.
  
      This returns a dictionary with keys:
  
      -  "sample_sample" map each downstream sample site to a list of the upstream
         sample sites on which it depend;
      -  "sample_dist" maps each sample site to the name of the distribution at
         that site;
      -  "plate_sample" maps each plate name to a list of the sample sites within
         that plate; and
      -  "observe" is a list of observed sample sites.
  
      For example for the model::
  
          def model(data):
              m = pyro.sample('m', dist.Normal(0, 1))
              sd = pyro.sample('sd', dist.LogNormal(m, 1))
              with pyro.plate('N', len(data)):
                  pyro.sample('obs', dist.Normal(m, sd), obs=data)
  
      the relation is::
  
          {'sample_sample': {'m': [], 'sd': ['m'], 'obs': ['m', 'sd']},
           'sample_dist': {'m': 'Normal', 'sd': 'LogNormal', 'obs': 'Normal'},
           'plate_sample': {'N': ['obs']},
           'observed': ['obs']}
  
      :param callable model: A model to inspect.
      :param model_args: Optional tuple of model args.
      :param model_kwargs: Optional dict of model kwargs.
      :param bool include_deterministic: Whether to include deterministic sites.
      :rtype: dict
      """
      if model_args is None:
          model_args = ()
      if model_kwargs is None:
          model_kwargs = {}
      assert isinstance(model_args, tuple)
      assert isinstance(model_kwargs, dict)
  
      with torch.random.fork_rng(), torch.no_grad(), pyro.validation_enabled(
          False
      ):
          with TrackProvenance(
              include_deterministic=include_deterministic
          ):
              trace = poutine.trace(model).get_trace(
                  *model_args, **model_kwargs
              )
  
      sample_sample = {}
      sample_param = {}
      sample_dist = {}
      param_constraint = {}
      plate_sample = defaultdict(list)
      observed = []
  
      def _get_type_from_frozenname(frozen_name):
          return trace.nodes[frozen_name]["type"]
  
      for name, site in trace.nodes.items():
          if site["type"] == "param":
              param_constraint[name] = str(site["kwargs"]["constraint"])
  
          if site["type"] != "sample" or site_is_subsample(site):
              continue
  
          provenance = get_provenance(
              site["fn"].log_prob(site["value"])
              if not site_is_deterministic(site)
              else site["fn"].base_dist.log_prob(site["value"])
          )
          sample_sample[name] = [
              upstream
              for upstream in provenance
              if upstream != name
              and _get_type_from_frozenname(upstream) == "sample"
          ]
  
          sample_param[name] = [
              upstream
              for upstream in get_provenance(
                  site["fn"].log_prob(site["value"])
              )
              if upstream != name
              and _get_type_from_frozenname(upstream) == "param"
          ]
  
          sample_dist[name] = (
              _get_dist_name(site["fn"])
              if not site_is_deterministic(site)
              else "Deterministic"
          )
          for frame in site["cond_indep_stack"]:
              plate_sample[frame.name].append(name)
          if site["is_observed"]:
              observed.append(name)
  
      def _resolve_plate_samples(plate_samples):
          for p, pv in plate_samples.items():
              pv = set(pv)
              for q, qv in plate_samples.items():
                  qv = set(qv)
                  if (
                      len(pv & qv) > 0
                      and len(pv - qv) > 0
                      and len(qv - pv) > 0
                  ):
                      plate_samples_ = plate_samples.copy()
                      plate_samples_[q] = pv & qv
                      plate_samples_[q + "__CLONE"] = qv - pv
                      return _resolve_plate_samples(plate_samples_)
          return plate_samples
  
      plate_sample = _resolve_plate_samples(plate_sample)
  
      # Normalize order of variables.
      def sort_by_time(names: Collection[str]) -> List[str]:
          return [name for name in trace.nodes if name in names]
  
      sample_sample = {
          k: sort_by_time(v) for k, v in sample_sample.items()
      }
      sample_param = {
          k: sort_by_time(v) for k, v in sample_param.items()
      }
      plate_sample = {
          k: sort_by_time(v) for k, v in plate_sample.items()
      }
      observed = sort_by_time(observed)
  
      return {
          "sample_sample": sample_sample,
          "sample_param": sample_param,
          "sample_dist": sample_dist,
          "param_constraint": param_constraint,
          "plate_sample": dict(plate_sample),
          "observed": observed,
      }
  
imports: import itertools
  import os
  from collections import defaultdict
  from pathlib import Path
  from types import SimpleNamespace
  from typing import Callable, Collection, Dict, List, Optional, Union
  import torch
  import pyro
  import pyro.distributions
  import pyro.poutine
  from pyro.ops.provenance import ProvenanceTensor, detach_provenance, get_provenance
  from pyro.poutine.messenger import Messenger
  from pyro.poutine.util import site_is_subsample
label_desc: Use image node_id 4 for calling a global function with example usage: get_model_relations(model, model_args, model_kwargs, include_deterministic) and returns: dict, trace, plate_samples, unknown, _resolve_plate_samples