Helper to get the data columns to plot.

def get_which_data_ycols(model, which_data_ycols):
    """
    Helper to get the data columns to plot.
    """
    if which_data_ycols == 'all' or which_data_ycols is None:
        return np.arange(model.output_dim)
    return which_data_ycols