code: def _tmc_diagonal_sample(msg: Message) -> torch.Tensor:
      assert isinstance(msg["fn"], TorchDistributionMixin)
      assert msg["infer"] is not None
      dist, num_samples = msg["fn"], msg["infer"]["num_samples"]
  
      # find batch dims that aren't plate dims
      batch_shape = [1] * len(dist.batch_shape)
      for f in msg["cond_indep_stack"]:
          if f.vectorized:
              assert f.dim is not None
              batch_shape[f.dim] = (
                  f.size if f.size > 0 else dist.batch_shape[f.dim]
              )
      batch_shape_tuple = tuple(batch_shape)
  
      # sample a batch
      sample_shape = (num_samples,)
      fat_sample = dist(
          sample_shape=torch.Size(sample_shape)
      )  # TODO thin before sampling
      assert (
          fat_sample.shape
          == sample_shape + dist.batch_shape + dist.event_shape
      )
      assert any(d > 1 for d in fat_sample.shape)
  
      target_shape = (
          (num_samples,) + batch_shape_tuple + dist.event_shape
      )
  
      # if this site has any ancestors, choose ancestors from diagonal approximation
      thin_sample = fat_sample
      if thin_sample.shape != target_shape:
          index: List[Any] = [Ellipsis] + [slice(None)] * (
              len(thin_sample.shape) - 1
          )
          squashed_dims = []
          for squashed_dim, squashed_size in zip(
              range(1, len(thin_sample.shape)), thin_sample.shape[1:]
          ):
              if squashed_size > 1 and (
                  target_shape[squashed_dim] == 1 or squashed_dim == 0
              ):
                  # diagonal approximation: identify particle indices across populations
                  ancestor_index = torch.arange(
                      squashed_size, device=thin_sample.device
                  )
                  index[squashed_dim] = ancestor_index
                  squashed_dims.append(squashed_dim)
  
          thin_sample = Vindex(thin_sample)[tuple(index)]
          for squashed_dim in squashed_dims:
              thin_sample = thin_sample.unsqueeze(squashed_dim)
  
      assert thin_sample.shape == target_shape
      return thin_sample
  
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 2 for calling a global function with example usage: _tmc_diagonal_sample(msg) and returns: thin_sample