def gather_dataset_statistics(grounded_scan):
  """Gather dataset statistics and save them for each split."""
  logging.info("Gathering dataset statistics...")

  splits = ["train", "test"]
  if FLAGS.split == "generalization":
    splits.extend([
        "visual", "situational_1", "situational_2", "contextual", "adverb_1",
        "adverb_2", "visual_easier"
    ])
  elif FLAGS.split == "spatial_relation":
    splits.extend([
        "visual", "relation", "referent", "relative_position_1",
        "relative_position_2"
    ])
  elif FLAGS.split == "target_lengths":
    splits.append("target_lengths")
  if FLAGS.make_dev_set:
    splits += ["dev"]
  for split in splits:
    grounded_scan.save_dataset_statistics(split=split)
    logging.info("Saved dataset staistics for split %s", split)
  if FLAGS.visualize_per_template > 0 or FLAGS.visualize_per_split > 0:
    grounded_scan.visualize_data_examples(
        visualize_per_split=FLAGS.visualize_per_split)
  if FLAGS.count_equivalent_examples:
    if FLAGS.split == "uniform":
      splits_to_count = ["test"]
    elif FLAGS.split == "generalization":
      splits_to_count = [
          "visual", "situational_1", "situational_2", "contextual"
      ]
    elif FLAGS.split == "spatial_relation":
      splits_to_count = [
          "visual", "relation", "referent", "relative_position_1",
          "relative_position_2"
      ]
    else:
      raise ValueError(f"Unknown option for flag --split: {FLAGS.split}")
    for split in splits_to_count:
      logging.info("Equivalent examples in train and %s: %s", split,
                   grounded_scan.count_equivalent_examples("train", split))
Gather dataset statistics and save them for each split.
language/gscan/data/main.py