def axis_off(ax, which='both'):
    """Turn off specific axis in an ``ax``."""
    assert which in ('x', 'y', 'both'), 'Which must be `x`, `y`, or `both`.'
    if which == 'both':
        ax.get_xaxis().set_visible(False)
        ax.get_yaxis().set_visible(False)
    if which == 'x':
        ax.get_xaxis().set_visible(False)
    elif which == 'y':
        ax.get_yaxis().set_visible(False)
    return ax
