def get_figure(ax):
    """
    Retrieves figure from axes. Axes can be either an instance
    of Matplotlib.Axes or a 1D/2D array of Matplotlib.Axes.

    :param ax: Axes or vector/array of Axes
    :return: Matplotlib.Figure
    """
    fig = None
    try:
        fig = ax.get_figure()
    except AttributeError:
        try:
            fig = ax[0].get_figure()
        except AttributeError:
            fig = ax[0][0].get_figure()
    return fig
