code: def select_samples(samples, num_samples=None, group_by_chain=False):
      """
      Performs selection from given MCMC samples.
  
      :param dictionary samples: Samples object to sample from.
      :param int num_samples: Number of samples to return. If `None`, all the samples
          from an MCMC chain are returned in their original ordering.
      :param bool group_by_chain: Whether to preserve the chain dimension. If True,
          all samples will have num_chains as the size of their leading dimension.
      :return: dictionary of samples keyed by site name.
      """
      if num_samples is None:
          # reshape to collapse chain dim when group_by_chain=False
          if not group_by_chain:
              samples = {
                  k: v.reshape((-1,) + v.shape[2:])
                  for k, v in samples.items()
              }
      else:
          if not samples:
              raise ValueError("No samples found from MCMC run.")
          if group_by_chain:
              batch_dim = 1
          else:
              samples = {
                  k: v.reshape((-1,) + v.shape[2:])
                  for k, v in samples.items()
              }
              batch_dim = 0
          sample_tensor = list(samples.values())[0]
          batch_size, device = (
              sample_tensor.shape[batch_dim],
              sample_tensor.device,
          )
          idxs = torch.randint(
              0, batch_size, size=(num_samples,), device=device
          )
          samples = {
              k: v.index_select(batch_dim, idxs)
              for k, v in samples.items()
          }
      return samples
  
imports: import functools
  import traceback
  import warnings
  from collections import OrderedDict, defaultdict
  from functools import partial, reduce
  from itertools import product
  import torch
  from opt_einsum import shared_intermediates
  from torch.distributions import biject_to
  import pyro
  import pyro.poutine
  from pyro.distributions.util import broadcast_shape, logsumexp
  from pyro.infer import config_enumerate
  from pyro.infer.autoguide.initialization import InitMessenger, init_to_uniform
  from pyro.infer.util import is_validation_enabled
  from pyro.ops import stats
  from pyro.ops.contract import contract_to_tensor
  from pyro.ops.integrator import potential_grad
  from pyro.poutine.subsample_messenger import _Subsample
  from pyro.poutine.util import prune_subsample_sites
  from pyro.util import check_site_shape, ignore_jit_warnings
label_desc: Use image node_id 10 for calling a global function with example usage: select_samples(samples, num_samples, group_by_chain) and returns: samples