def _pack(self, dataset, window_size=None, cycle_length=None,
            deterministic=False, keys=None):
    """Main method for chaining together packing transformation steps."""
    (dataset, self._num_sequences, self._token_dtype, keys
    ) = self._standardize(dataset, keys)
    if window_size is None:
      dataset = self._scanning_pack(dataset)
    else:
      # Dataset.window splits nested Tensors.
      re_zip = lambda *x: tf.data.Dataset.zip(x)
      dataset = dataset.window(window_size).map(re_zip).interleave(
          self._scanning_pack, cycle_length=cycle_length,
          block_length=window_size,
          num_parallel_calls=tf.data.experimental.AUTOTUNE)

      if not deterministic:
        # Sloppy interleave offers a marginal performance improvement.
        options = tf.data.Options()
        options.experimental_deterministic = False
        dataset = dataset.with_options(options)

    dataset = dataset.map(
        self._finalize, num_parallel_calls=tf.data.experimental.AUTOTUNE)
    self._num_sequences, self._token_dtype = None, None

    if keys:
      def dict_pack(example):
        output = {}
        for i, key in enumerate(keys):
          output[key] = example["contents"][:, i]
          output[key + "_segmentation"] = example["segment"][:, i]
          output[key + "_position"] = example["position"][:, i]
        return output
      dataset = dataset.map(dict_pack)
    return dataset
Main method for chaining together packing transformation steps.
tensor2tensor/data_generators/generator_utils.py