def _get_birds200_dataset(
    mode: str,
    rng: np.ndarray) -> Tuple[tf.data.Dataset, tf.data.Dataset, int]:
  """Load the caltech_birds2011 dataset."""
  assert jax.host_count() == 1, (
      "caltech_birds2011 dataset does not support multihost training. "
      "Found {} hosts.".format(jax.host_count()))

  dataset_builder = tfds.builder("caltech_birds2011")
  num_classes = 200

  # Make sure each host uses a different RNG for the training data.
  rng, data_rng = jax.random.split(rng)
  data_rng = jax.random.fold_in(data_rng, jax.host_id())
  data_rng, shuffle_rng = jax.random.split(data_rng)

  if mode == "train-val":
    read_config = tfds.ReadConfig(shuffle_seed=shuffle_rng[0])
    ds = dataset_builder.as_dataset(
        split="train", shuffle_files=False, read_config=read_config)

    train_ds = ds.take(5000).shuffle(5000, seed=shuffle_rng[0])
    eval_ds = ds.skip(5000)

  elif mode == "train-test":
    train_split = "train"
    eval_split = "test"

    train_read_config = tfds.ReadConfig(shuffle_seed=shuffle_rng[0])
    train_ds = dataset_builder.as_dataset(
        split=train_split, shuffle_files=True, read_config=train_read_config)

    eval_read_config = tfds.ReadConfig(shuffle_seed=shuffle_rng[1])
    eval_ds = dataset_builder.as_dataset(
        split=eval_split, shuffle_files=False, read_config=eval_read_config)
  else:
    raise ValueError(f"Unknown mode: {mode}.")

  return train_ds, eval_ds, num_classes