def plot_k_fold(ax, cv, all_dataset, X_train, y_train, test_set):
    """Create a sample plot for training, validation, testing."""
    bd_cmap = [
        "#3A6190",
        "#683E00",
        "#2196F3",
        "#DDF2FF",
    ]

    ax.barh("Original\nDataset", len(all_dataset), left=0, height=0.5, color=bd_cmap[0])

    # Generate the training/validation/testing data fraction visualizations for each CV split
    for ii, (tr_idx, val_idx) in enumerate(cv.split(X=X_train, y=y_train)):
        n_train, n_val, n_test = len(tr_idx), len(val_idx), len(test_set)
        n_train2 = n_train + n_val - max(val_idx) - 1
        ax.barh("cv" + str(ii + 1), min(val_idx), left=0, height=0.5, color=bd_cmap[1])
        ax.barh(
            "cv" + str(ii + 1), n_val, left=min(val_idx), height=0.5, color=bd_cmap[2]
        )
        ax.barh(
            "cv" + str(ii + 1),
            n_train2,
            left=max(val_idx) + 1,
            height=0.5,
            color=bd_cmap[1],
        )
        ax.barh(
            "cv" + str(ii + 1),
            n_test,
            left=n_train + n_val,
            height=0.5,
            color=bd_cmap[3],
        )

    ax.invert_yaxis()
    ax.set_xlim([-int(0.1 * len(all_dataset)), int(1.1 * len(all_dataset))])
    ax.set(xlabel="Number of samples.", title="KFold Train-Test-Valid split")
    ax.legend(
        [Patch(color=bd_cmap[i]) for i in range(4)],
        ["Original set", "Training set", "Validation set", "Testing set"],
        loc="lower center",
        ncols=2,
    )
    ax.text(
        -0.07,
        0.45,
        "Train-Valid-Test split",
        rotation=90,
        verticalalignment="center",
        horizontalalignment="left",
        transform=ax.transAxes,
    )
    return ax
Create a sample plot for training, validation, testing.
examples/model_building/plot_how_train_test_and_tune.py