docs: def train_model(self, balance, *args, **kwargs):
          """
          Args:
              balance: A 1d arraylike that sums to 1, corresponding to the
                  (possibly estimated) class balance.
          """
          self.balance = np.array(balance)