def check_treedepth(samples, max_treedepth=10, quiet=False,
    return_diagnostics=False):
    """Check transitions that ended prematurely due to maximum tree depth limit.

    Parameters
    ----------
    samples : ArviZ InferenceData instance
        Contains samples to be checked. Must contain both `posterior`
        and `samples_stats`.
    max_treedepth: int, default 10
        Maximum tree depth used in the calculation.
    quiet : bool, default False
        If True, do not print diagnostic result to the screen.
    return_diagnostics : bool, default False
        If True, return both a Boolean about whether the diagnostic
        passed and the number of samples where the tree depth was too
        deep. Otherwise, only return Boolean if the test passed.

    Returns
    -------
    passed : bool
        Return True if tree depth test passed. Return False otherwise.
    n_too_deep : int, optional
        Number of samplers wherein the tree depth was greater than
        `max_treedepth`.
    """
    n_too_deep = (samples.sample_stats.treedepth.values >= max_treedepth).sum()
    n_total = samples.sample_stats.dims['chain'] * samples.sample_stats.dims[
        'draw']
    if not quiet:
        msg = '{} of {} ({}%) iterations saturated'.format(n_too_deep,
            n_total, 100 * n_too_deep / n_total)
        msg += ' the maximum tree depth of {}.'.format(max_treedepth)
        print(msg)
    pass_check = n_too_deep == 0
    if not pass_check and not quiet:
        print(
            '  Try running again with max_treedepth set to a larger value' +
            ' to avoid saturation.')
    if return_diagnostics:
        return pass_check, n_too_deep
    return pass_check
