def multibootstrap(run_info,
                   preds,
                   labels,
                   score_fn,
                   paired_seeds = False,
                   sample_seeds = True,
                   sample_examples = True,
                   nboot = 500,
                   rng = np.random,
                   verbose = True,
                   progress_indicator=lambda x: x):
  """Compute bootstrap samples of score_fn using multibootstrap over seeds and examples.

  Args:
    run_info: DataFrame with columns 'seed' (any type) and 'intervention' (bool)
      Any other columns will be ignored, and all rows with the same value of
      (seed, intervention) will have their scores averaged. This is done as a
      macro-average, such that each seed contributes equally to the final
      average.
    preds: [num_runs, num_examples] matrix of predictions, aligned to rows of
      run_info.
    labels: [num_examples] array of target labels, aligned to columns of preds.
    score_fn: function that takes (y_true: List, y_pred: List) and returns a
      scalar.
    paired_seeds: if True, assumes base and expt are paired on the same set of
      seeds, and will use the same set of sampled seeds for both sides. If
      False, will sample seeds independently but will still use paired sets of
      examples.
    sample_seeds: if True (default), bootstrap over seeds.
    sample_examples: if True (default), bootstrap over examples.
    nboot: number of bootstrap samples to compute.
    rng: random number generator state, or an int to seed a new one.
    verbose: set to False to disable printing of intermediate info.
    progress_indicator: optional, pass-through iterator like tqdm to indicate
      progress, as computing many samples can take a while.

  Returns:
    List[Tuple[float, float]] of nboot paired bootstrap samples (L, L')
  """
  assert len(run_info) == len(preds)
  assert len(labels) == preds.shape[1]
  num_examples = len(labels)
  # Reset dataframe index to be row indices [0, 1, ..., num_runs]
  run_info = run_info.reset_index(drop=True)
  # Map seed -> list[int] of row indices into preds
  mask = run_info['intervention']
  base_seed_to_row = run_info[~mask].groupby('seed').groups
  expt_seed_to_row = run_info[mask].groupby('seed').groups

  if isinstance(rng, int):
    rng = np.random.RandomState(rng)

  if paired_seeds:
    if verbose:
      print(f'Multibootstrap (paired) on {num_examples} examples')
    # Get set of seed keys in intersection
    common_seeds = list(
        set(base_seed_to_row.keys()).intersection(expt_seed_to_row.keys()))
    if verbose:
      print(f'  Common seeds ({len(common_seeds)}): {str(common_seeds)}')
      print(f'  Base: {len(run_info[~mask])} runs')
      print(f'  Expt: {len(run_info[mask])} runs')
    # Warn if there are any seeds which don't overlap
    base_only_seeds = set(base_seed_to_row.keys()).difference(
        expt_seed_to_row.keys())
    if verbose and base_only_seeds:
      print(f'  Warning: {len(base_only_seeds)} seeds found on base side only. '
            f'These will be omitted from the analysis: {str(base_only_seeds)}')
    expt_only_seeds = set(expt_seed_to_row.keys()).difference(
        base_seed_to_row.keys())
    if verbose and expt_only_seeds:
      print(f'  Warning: {len(expt_only_seeds)} seeds found on expt side only. '
            f'These will be omitted from the analysis: {str(expt_only_seeds)}')
  else:
    base_seeds = list(sorted(base_seed_to_row.keys()))
    expt_seeds = list(sorted(expt_seed_to_row.keys()))
    if verbose:
      print(f'Multibootstrap (unpaired) on {num_examples} examples')
      print(f'  Base seeds ({len(base_seeds)}): {str(base_seeds)}')
      print(f'  Base: {len(run_info[~mask])} runs')
      print(f'  Expt seeds ({len(expt_seeds)}): {str(expt_seeds)}')
      print(f'  Expt: {len(run_info[mask])} runs')

  def _score_seeds(seeds_to_rows, seeds,
                   row_score_fn):
    """Helper to score a single bootstrap sample."""
    scores = []
    for s in seeds:
      ii = seeds_to_rows[s]
      seed_scores = [row_score_fn(i) for i in ii]
      scores.append(np.mean(seed_scores))  # mean over finetuning params, etc.
    return np.mean(scores)  # mean over seeds

  # Compute bootstrap samples
  bootstrap_samples = []
  for _ in progress_indicator(range(nboot)):
    # Sample seeds; we'll map these to row indices later.
    if sample_seeds:
      if paired_seeds:
        base_seed_sample = _bsample(common_seeds, rng)
        expt_seed_sample = base_seed_sample
      else:
        base_seed_sample = _bsample(base_seeds, rng)
        expt_seed_sample = _bsample(expt_seeds, rng)
    else:
      if paired_seeds:
        base_seed_sample = common_seeds
        expt_seed_sample = common_seeds
      else:
        base_seed_sample = base_seeds
        expt_seed_sample = expt_seeds

    # Sample examples and select columns
    if sample_examples:
      jj = rng.choice(num_examples, size=num_examples, replace=True)
      selected_labels = labels[jj]
      # pylint note: safe to close over jj, because this function is also
      # redefined on every loop iteration.
      score_row = lambda i: score_fn(selected_labels, preds[i, jj])  # pylint: disable=cell-var-from-loop
    else:
      score_row = lambda i: score_fn(labels, preds[i])  # pylint: disable=cell-var-from-loop  # pytype: disable=wrong-arg-types

    # Select and score rows corresponding to selected seeds.
    base_score = _score_seeds(base_seed_to_row, base_seed_sample, score_row)
    expt_score = _score_seeds(expt_seed_to_row, expt_seed_sample, score_row)

    # Store this sample as a pair.
    bootstrap_samples.append((base_score, expt_score))

  return bootstrap_samples
Compute bootstrap samples of score_fn using multibootstrap over seeds and examples.

  Args:
    run_info: DataFrame with columns 'seed' (any type) and 'intervention' (bool)
      Any other columns will be ignored, and all rows with the same value of
      (seed, intervention) will have their scores averaged. This is done as a
      macro-average, such that each seed contributes equally to the final
      average.
    preds: [num_runs, num_examples] matrix of predictions, aligned to rows of
      run_info.
    labels: [num_examples] array of target labels, aligned to columns of preds.
    score_fn: function that takes (y_true: List, y_pred: List) and returns a
      scalar.
    paired_seeds: if True, assumes base and expt are paired on the same set of
      seeds, and will use the same set of sampled seeds for both sides. If
      False, will sample seeds independently but will still use paired sets of
      examples.
    sample_seeds: if True (default), bootstrap over seeds.
    sample_examples: if True (default), bootstrap over examples.
    nboot: number of bootstrap samples to compute.
    rng: random number generator state, or an int to seed a new one.
    verbose: set to False to disable printing of intermediate info.
    progress_indicator: optional, pass-through iterator like tqdm to indicate
      progress, as computing many samples can take a while.

  Returns:
    List[Tuple[float, float]] of nboot paired bootstrap samples (L, L')
  
language/multiberts/multibootstrap.py