def generate_batch_spec(mode, batch_size):
    """ Generates a spec describing how to draw batches

    Args:
      mode: one of ['train', 'test', 'val']
    """
    assert mode in ['train', 'test', 'val']

    # on a more complicated dataset this would include useful arguments
    # such as whether to augment the data, what data source to draw from....
    # this example is simple enough to not need a spec so this is just for illustration purposes
    batch_spec = {
        'mode': mode,
        'batch_size': batch_size
    }
    return batch_spec