code: def plot_within_chain(
      model: SklearnModel,
      ax=None,
      title="Within Chain Variation",
      x_label=False,
      X=None,
      y=None,
  ):
      if ax is None:
          _, ax = plt.subplots(1, 1)
      n_chains = model.n_chains
  
      chain_len = int(len(model.model_samples) / n_chains)
      color = iter(cm.rainbow(np.linspace(0, 1, n_chains)))
  
      for c in range(n_chains):
          clr = next(color)
          chain_preds = model.chain_predictions(X, c)
          mean_pred = np.array(chain_preds).mean(axis=0)
  
          y_plt = [
              mean_squared_error(mean_pred, p) for p in chain_preds
          ]
          ax.plot(
              np.arange(chain_len),
              y_plt,
              color=clr,
              label=f"Chain {c} (Average {np.round(np.mean(y_plt), 2)})",
          )
  
      ax.set_ylabel("mean squared distance to average iteration")
      if x_label:
          ax.set_xlabel("Iteration")
      ax.legend()
      ax.set_title(title)
      return ax
  
imports: import argparse
  import itertools
  import os
  from functools import partial
  import numpy
  import pandas
  import matplotlib.pyplot
  from matplotlib import cm
  from tqdm import tqdm
  from sklearn import model_selection, datasets
  from sklearn.metrics import mean_squared_error
  from imodels import get_clean_dataset
  from imodels.experimental.bartpy.model import Model
  from imodels.experimental.bartpy.tree import Tree
  from ..sklearnmodel import BART, SklearnModel
label_desc: Use image node_id 6 for calling a global function with example usage: plot_within_chain(model, ax, title, x_label, X, y) and returns: ax