code: def enum_extend(
      trace: "Trace", msg: "Message", num_samples: Optional[int] = None
  ) -> List["Trace"]:
      """
      :param trace: a partial trace
      :param msg: the message at a Pyro primitive site
      :param num_samples: maximum number of extended traces to return.
      :returns: a list of traces, copies of input trace with one extra site
  
      Utility function to copy and extend a trace with sites based on the input site
      whose values are enumerated from the support of the input site's distribution.
  
      Used for exact inference and integrating out discrete variables.
      """
      if num_samples is None:
          num_samples = -1
  
      extended_traces = []
      assert msg["name"] is not None
      if TYPE_CHECKING:
          assert isinstance(msg["fn"], Distribution)
      for i, s in enumerate(
          msg["fn"].enumerate_support(*msg["args"], **msg["kwargs"])
      ):
          if i > num_samples and num_samples >= 0:
              break
          msg_copy = msg.copy()
          msg_copy.update(value=s)  # type: ignore[call-arg]
          tr_cp = trace.copy()
          tr_cp.add_node(msg["name"], **msg_copy)
          extended_traces.append(tr_cp)
      return extended_traces
  
file_path: pyro/pyro/poutine/util.py
imports: from typing import TYPE_CHECKING, List, Optional
  from pyro import settings
label_desc: Use image node_id 6 for calling a global function with example usage: enum_extend(trace, msg, num_samples) and returns: extended_traces