def _sample_distribution_over_matrix(rng):
  """Samples the config for a distribution over a matrix.

  Args:
    rng: np.random.RandomState

  Returns:
    A distribution over matrix config (containing a tuple with the name and
    extra args needed to create the distribution).
  """
  # At this point, the choices here are arbitrary. We know some things about
  # the spectrum of the hessian for various problems we care about but certainly
  # not all. For example see: https://arxiv.org/abs/1901.10159.
  # These selections define a distribution over spectrum. This distribution is
  # NOT like the ones found in prior work, but tries to be broader, and
  # (hopefully) will contain spectrum from real problems.
  choice = rng.choice(["normal", "linspace_eigen", "uniform", "logspace_eigen"])
  if choice == "normal":
    mean = rng.normal(0, 1)
    std = rng.uniform(0, 2)
    return choice, {"mean": mean, "std": std}
  elif choice == "uniform":
    min_v = rng.uniform(-3, 1)
    max_v = min_v + rng.uniform(0, 5)
    return choice, {"min": min_v, "max": max_v}
  elif choice == "linspace_eigen":
    min_v = rng.uniform(0.01, 50)
    max_v = min_v + rng.uniform(0, 100)
    return choice, {"min": min_v, "max": max_v}
  elif choice == "logspace_eigen":
    min_v = 1e-5
    max_v = rng.uniform(1, 1000)
    return choice, {"min": min_v, "max": max_v}