code: def enumerate_site(msg: Message) -> torch.Tensor:
      assert isinstance(msg["fn"], TorchDistributionMixin)
      assert msg["infer"] is not None
      dist, num_samples = msg["fn"], msg["infer"].get("num_samples")
      if num_samples is None:
          # Enumerate over the support of the distribution.
          value = dist.enumerate_support(
              expand=msg["infer"].get("expand", False)
          )
      elif num_samples > 1 and not msg["infer"].get("expand", False):
          tmc_strategy = msg["infer"].get("tmc", "diagonal")
          if tmc_strategy == "mixture":
              value = _tmc_mixture_sample(msg)
          elif tmc_strategy == "diagonal":
              value = _tmc_diagonal_sample(msg)
          else:
              raise ValueError(
                  "{} not a valid TMC strategy".format(tmc_strategy)
              )
      elif num_samples > 1 and msg["infer"]["expand"]:
          # Monte Carlo sample the distribution.
          value = dist(sample_shape=torch.Size([num_samples]))
      assert value.dim() == 1 + len(dist.batch_shape) + len(
          dist.event_shape
      )
      return value
  
file_path: pyro/pyro/poutine/enum_messenger.py
imports: from typing import Any, Dict, List, Optional
  import torch
  from typing_extensions import Self
  from pyro.distributions.torch import Categorical
  from pyro.distributions.torch_distribution import TorchDistributionMixin
  from pyro.ops.indexing import Vindex
  from pyro.poutine.messenger import Messenger
  from pyro.poutine.runtime import _ENUM_ALLOCATOR, Message
  from pyro.util import ignore_jit_warnings
label_desc: Use image node_id 3 for calling a global function with example usage: enumerate_site(msg) and returns: value