def _standardize(self, dataset, keys):
    """Force dataset structure into a tuple of Tensors."""
    shapes = tf.data.get_output_shapes(dataset)

    if isinstance(shapes, dict):
      keys = keys or tuple(shapes.keys())
      dataset = dataset.map(lambda x: tuple(x[k] for k in keys))
      shapes = tf.data.get_output_shapes(dataset)

    if not all(isinstance(i, tf.TensorShape) for i in shapes):
      # Internally this class expects tuples of Tensors, even for the degenerate
      # case of a single sequence.
      dataset = dataset.map(lambda x: (x,))
      shapes = tf.data.get_output_shapes(dataset)

    for s in shapes:
      if not s.is_compatible_with(tf.TensorShape([None])):
        raise ValueError("Tensors to be packed must be one-dimensional.")

    if not shapes:
      raise ValueError("Expected sequence dataset.")

    if self._chop_long_sequences and len(shapes) != 1:
      raise ValueError("chop_long_sequences expects a single sequence dataset.")

    token_types = tf.data.get_output_types(dataset)
    if len(set(token_types)) > 1:
      raise ValueError("Inconsistent dtypes: {}".format(token_types))

    return dataset, len(shapes), token_types[0], keys
Force dataset structure into a tuple of Tensors.
tensor2tensor/data_generators/generator_utils.py