def METHOD_NAME(tag):
    """Return the the number of columns to use in tests."""
    n_columns_list = []
    if tag in ["univariate", "both"]:
        n_columns_list = [1, 2]
    elif tag == "multivariate":
        n_columns_list = [2]
    else:
        raise ValueError(f"Unexpected tag {tag} in _get_n_columns.")
    return n_columns_list