code: def class_balance(y_train, y_test=None, ax=None, labels=None, **kwargs):
      
      # Instantiate the visualizer
      visualizer = ClassBalance(ax=ax, labels=labels, **kwargs)
  
      # Fit and transform the visualizer (calls draw)
      visualizer.fit(y_train, y_test)
      visualizer.finalize()
  
      # Return the axes object on the visualizer
      return visualizer.ax
docs: Quick method:
  
      One of the biggest challenges for classification models is an imbalance of
      classes in the training data. This function vizualizes the relationship of
      the support for each class in both the training and test data by
      displaying how frequently each class occurs as a bar graph.
  
      The figure can be displayed in two modes:
  
      1. Balance mode: show the frequency of each class in the dataset.
      2. Compare mode: show the relationship of support in train and test data.
  
      Balance mode is the default if only y_train is specified. Compare mode
      happens when both y_train and y_test are specified.
  
      Parameters
      ----------
      y_train : array-like
          Array or list of shape (n,) that containes discrete data.
  
      y_test : array-like, optional
          Array or list of shape (m,) that contains discrete data. If
          specified, the bar chart will be drawn in compare mode.
  
      ax : matplotlib Axes, default: None
          The axis to plot the figure on. If None is passed in the current axes
          will be used (or generated if required).
  
      labels: list, optional
          A list of class names for the x-axis if the target is already encoded.
          Ensure that the labels are ordered lexicographically with respect to
          the values in the target. A common use case is to pass
          LabelEncoder.classes\_ as this parameter. If not specified, the labels
          in the data will be used.
  
      kwargs: dict, optional
          Keyword arguments passed to the super class. Here, used
          to colorize the bars in the histogram.
  
      Returns
      -------
      ax : matplotlib axes
          Returns the axes that the class balance plot was drawn on.