docs: def plot_chain(sampler, p=None, **kwargs):
      """Generate a diagnostic plot of the sampler chains.
  
      Parameters
      ----------
      sampler : `emcee.EnsembleSampler`
          Sampler containing the chains to be plotted.
      p : int (optional)
          Index of the parameter to plot. If omitted, all chains are plotted.
      last_step : bool (optional)
          Whether to plot the last step of the chain or the complete chain
          (default).
  
      Returns
      -------
      figure : `matplotlib.figure.Figure`
          Figure
      """
      if p is None:
          npars = sampler.chain.shape[-1]
          for pp in six.moves.range(npars):
              _plot_chain_func(sampler, pp, **kwargs)
          fig = None
      else:
          fig = _plot_chain_func(sampler, p, **kwargs)
  
      return fig