code: def _sample_return(
      *,
      run: Optional[RunType],
      traces: Sequence[IBaseTrace],
      tune: int,
      t_sampling: float,
      discard_tuned_samples: bool,
      compute_convergence_checks: bool,
      return_inferencedata: bool,
      keep_warning_stat: bool,
      idata_kwargs: Dict[str, Any],
      model: Model,
  ) -> Union[InferenceData, MultiTrace]:
      """Final step of `pm.sampler` that picks/slices chains,
      runs diagnostics and converts to the desired return type."""
      # Pick and slice chains to keep the maximum number of samples
      if discard_tuned_samples:
          traces, length = _choose_chains(traces, tune)
      else:
          traces, length = _choose_chains(traces, 0)
      mtrace = MultiTrace(traces)[:length]
  
      # count the number of tune/draw iterations that happened
      # ideally via the "tune" statistic, but not all samplers record it!
      if "tune" in mtrace.stat_names:
          # Get the tune stat directly from chain 0, sampler 0
          stat = mtrace._straces[0].get_sampler_stats(
              "tune", sampler_idx=0
          )
          stat = tuple(stat)
          n_tune = stat.count(True)
          n_draws = stat.count(False)
      else:
          # these may be wrong when KeyboardInterrupt happened, but they're better than nothing
          n_tune = min(tune, len(mtrace))
          n_draws = max(0, len(mtrace) - n_tune)
  
      if discard_tuned_samples:
          mtrace = mtrace[n_tune:]
  
      # save metadata in SamplerReport
      mtrace.report._n_tune = n_tune
      mtrace.report._n_draws = n_draws
      mtrace.report._t_sampling = t_sampling
  
      n_chains = len(mtrace.chains)
      _log.info(
          f'Sampling {n_chains} chain{"s" if n_chains > 1 else ""} for {n_tune:_d} tune and {n_draws:_d} draw iterations '
          f"({n_tune*n_chains:_d} + {n_draws*n_chains:_d} draws total) "
          f"took {t_sampling:.0f} seconds."
      )
  
      idata = None
      if compute_convergence_checks or return_inferencedata:
          ikwargs: Dict[str, Any] = dict(
              model=model, save_warmup=not discard_tuned_samples
          )
          ikwargs.update(idata_kwargs)
          idata = pm.to_inference_data(mtrace, **ikwargs)
  
          if compute_convergence_checks:
              warns = run_convergence_checks(idata, model)
              mtrace.report._add_warnings(warns)
              log_warnings(warns)
  
          if return_inferencedata:
              # By default we drop the "warning" stat which contains `SamplerWarning`
              # objects that can not be stored with `.to_netcdf()`.
              if not keep_warning_stat:
                  return drop_warning_stat(idata)
              return idata
      return mtrace
  
imports: import logging
  import pickle
  import sys
  import time
  import warnings
  from typing import Any, Dict, Iterator, List, Literal, Mapping, Optional, Sequence, Set, Tuple, Type, Union, overload
  import numpy
  import pytensor.gradient
  from arviz import InferenceData, dict_to_dataset
  from arviz.data.base import make_attrs
  from fastprogress.fastprogress import progress_bar
  from pytensor.graph.basic import Variable
  from typing_extensions import Protocol, TypeAlias
  import pymc
  from pymc.backends import RunType, TraceOrBackend, init_traces
  from pymc.backends.arviz import coords_and_dims_for_inferencedata, find_constants, find_observations
  from pymc.backends.base import IBaseTrace, MultiTrace, _choose_chains
  from pymc.blocking import DictToArrayBijection
  from pymc.exceptions import SamplingError
  from pymc.initial_point import PointType, StartDict, make_initial_point_fns_per_chain
  from pymc.model import Model, modelcontext
  from pymc.sampling.parallel import Draw, _cpu_count
  from pymc.sampling.population import _sample_population
  from pymc.stats.convergence import log_warning_stats, log_warnings, run_convergence_checks
  from pymc.step_methods import NUTS, CompoundStep
  from pymc.step_methods.arraystep import BlockedStep, PopulationArrayStepShared
  from pymc.step_methods.hmc import quadpotential
  from pymc.util import RandomSeed, RandomState, _get_seeds_per_chain, drop_warning_stat, get_untransformed_name, is_transformed_name
  from pymc.vartypes import discrete_types
label_desc: Use image node_id 9 for calling a global function with example usage: _sample_return() and returns: mtrace, idata, drop_warning_stat