def METHOD_NAME(dataset):
  """Converts `dataset` for use with the output of `create_simple_keras_model`.
  Args:
    dataset: An instance of `tf.data.Dataset` to read from.
  Returns:
    An instance of `tf.data.Dataset` after conversion.
  """
  def map_fn(example):
    return collections.OrderedDict([
        ('x', tf.reshape(example['pixels'], [-1])),
        ('y', example['label']),
    ])
  return dataset.map(map_fn)