docs: def get_data(name, train_batch_size, test_batch_size):
    """Gets training and testing dataset iterators.
  
    Args:
      name: String. Name of dataset, either 'mnist' or 'cifar10'.
      train_batch_size: Integer. Batch size for training.
      test_batch_size: Integer. Batch size for testing.
  
    Returns:
      Dict containing:
        train_iterator: A tf.data.Iterator, over training data.
        test_iterator: A tf.data.Iterator, over test data.
        num_classes: Integer. Number of class labels.
    """
    if name not in ['mnist', 'cifar10']:
      raise ValueError(
          'Expected dataset \'mnist\' or \'cifar10\', but got %s' % name)
    dataset = getattr(tf.keras.datasets, name)
    num_classes = 10
  
    # Extract the raw data.
    raw_data = dataset.load_data()
    (images_train, labels_train), (images_test, labels_test) = raw_data
  
    # Normalize inputs and fix types.
    images_train = images_train.astype(np.float32) / 255.
    images_test = images_test.astype(np.float32) / 255.
    labels_train = labels_train.astype(np.int32).squeeze()
    labels_test = labels_test.astype(np.int32).squeeze()
  
    # Add a dummy 'color channel' dimension if it is not present.
    if images_train.ndim == 3:
      images_train = np.expand_dims(images_train, -1)
      images_test = np.expand_dims(images_test, -1)
  
    # Put the data onto the graph as constants.
    train_data = tf.data.Dataset.from_tensor_slices((images_train, labels_train))
    test_data = tf.data.Dataset.from_tensor_slices((images_test, labels_test))
  
    # Create iterators for each dataset.
    train_iterator = (
        train_data
        # Note: For larger datasets e.g. ImageNet, it will not be feasible to have
        # a shuffle buffer this large.
        .shuffle(buffer_size=len(images_train))
        .batch(train_batch_size)
        .repeat()
        .make_one_shot_iterator()
    )
    test_iterator = test_data.batch(test_batch_size).make_initializable_iterator()
    return dict(
        train_iterator=train_iterator,
        test_iterator=test_iterator,
        num_classes=num_classes)