class_docstr: Return a handler that records the inputs and outputs of primitive calls
  and their dependencies.
  
  Consider the following Pyro program:
  
      >>> def model(x):
      ...     s = pyro.param("s", torch.tensor(0.5))
      ...     z = pyro.sample("z", dist.Normal(x, s))
      ...     return z ** 2
  
  We can record its execution using ``trace``
  and use the resulting data structure to compute the log-joint probability
  of all of the sample sites in the execution or extract all parameters.
  
      >>> trace = pyro.poutine.trace(model).get_trace(0.0)
      >>> logp = trace.log_prob_sum()
      >>> params = [trace.nodes[name]["value"].unconstrained() for name in trace.param_nodes]
  
  :param fn: a stochastic function (callable containing Pyro primitive calls)
  :param graph_type: string that specifies the kind of graph to construct
  :param param_only: if true, only records params and not samples
  :returns: stochastic function decorated with a :class:`~pyro.poutine.trace_messenger.TraceMessenger`
docstr: :param string graph_type: string that specifies the type of graph
      to construct (currently only "flat" or "dense" supported)
  :param param_only: boolean that specifies whether to record sample sites
docstr_tok: [":","param","string","graph_type",":","string","that","specifies","the","type","of","graph","to","construct","(","currently","only","``","flat","''","or","``","dense","''","supported",")",":","param","param_only",":","boolean","that","specifies","whether","to","record","sample","sites"]
code: def __init__(
      self,
      graph_type: Optional[Literal["flat", "dense"]] = None,
      param_only: Optional[bool] = None,
  ) -> None:
      """
      :param string graph_type: string that specifies the type of graph
          to construct (currently only "flat" or "dense" supported)
      :param param_only: boolean that specifies whether to record sample sites
      """
      super().__init__()
      if graph_type is None:
          graph_type = "flat"
      if param_only is None:
          param_only = False
      assert graph_type in ("flat", "dense")
      self.graph_type = graph_type
      self.param_only = param_only
      self.trace = Trace(graph_type=self.graph_type)
  
imports: import sys
  from typing import TYPE_CHECKING, Any, Callable, Literal, Optional
  from typing_extensions import Self
  from pyro.poutine.messenger import Messenger
  from pyro.poutine.trace_struct import Trace
  from pyro.poutine.util import site_is_subsample
label_desc: Use image node_id 1 to create a new TraceMessenger object from inherited base classes: Messenger with example: obj = TraceMessenger(graph_type, param_only)