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.

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)