code: def align_samples(samples, model, particle_dim):
      """
      Unsqueeze stacked samples such that their particle dim all aligns.
      This traces ``model`` to determine the ``event_dim`` of each site.
      """
      assert particle_dim < 0
  
      sample = {name: value[0] for name, value in samples.items()}
      with poutine.block(), poutine.trace() as tr, poutine.condition(
          data=sample
      ):
          model()
  
      samples = samples.copy()
      for name, value in samples.items():
          event_dim = tr.trace.nodes[name]["fn"].event_dim
          pad = event_dim - particle_dim - value.dim()
          if pad < 0:
              raise ValueError(
                  "Cannot align samples, try moving particle_dim left"
              )
          if pad > 0:
              shape = value.shape[:1] + (1,) * pad + value.shape[1:]
              samples[name] = value.reshape(shape)
  
      return samples
  
file_path: pyro/pyro/contrib/epidemiology/util.py
imports: import numpy
  import torch
  import pyro
  import pyro.distributions
  import pyro.poutine
  from pyro.distributions.util import broadcast_shape
  from pyro.ops.special import safe_log
label_desc: Use image node_id 3 for calling a global function with example usage: align_samples(samples, model, particle_dim) and returns: samples