docs: def fit(self, y_train, y_test=None):
          """
          Fit the visualizer to the the target variables, which must be 1D
          vectors containing discrete (classification) data. Fit has two modes:
  
          1. Balance mode: if only y_train is specified
          2. Compare mode: if both train and test are specified
  
          In balance mode, the bar chart is displayed with each class as its own
          color. In compare mode, a side-by-side bar chart is displayed colored
          by train or test respectively.
  
          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.
          """
  
          # check to make sure that y_train is not a 2D array, e.g. X
          if y_train.ndim == 2:
              raise YellowbrickValueError((
                  "fit has changed to only require a 1D array, y "
                  "since version 0.9; please see the docs for more info"
              ))
  
          # Check the target types for the y variables
          self._validate_target(y_train)
          self._validate_target(y_test)
  
          # Get the unique values from the dataset
          targets = (y_train,) if y_test is None else (y_train, y_test)
          self.classes_ = unique_labels(*targets)
  
          # Validate the classes with the class names
          if self.labels is not None:
              if len(self.labels) != len(self.classes_):
                  raise YellowbrickValueError((
                      "discovered {} classes in the data, does not match "
                      "the {} labels specified."
                  ).format(len(self.classes_), len(self.labels)))
  
          # Determine if we're in compare or balance mode
          self._mode = BALANCE if y_test is None else COMPARE
  
          # Compute the support values
          if self._mode == BALANCE:
              self.support_ = np.array([
                  (y_train == idx).sum() for idx in self.classes_
              ])
  
          else:
              self.support_ = np.array([
                  [
                      (y == idx).sum() for idx in self.classes_
                  ]
                  for y in targets
              ])
  
          # Draw the bar chart
          self.draw()
  
          # Fit returns self
          return self