code: def pack_dataset(dataset, length, keys=None, use_custom_ops=False):
      """Creates a 'packed' version of a dataset on-the-fly.
  
      Borrowed from the tensor2tensor library.
      TODO(noam): make this faster
  
      This is meant to replace the irritation of having to create a separate
      "packed" version of a dataset to train efficiently on TPU.
  
      Each example in the output dataset represents several examples in the
      input dataset.
  
      For each key in the input dataset, two additional keys are created:
      <key>_segmentation: an int32 tensor identifying the parts
         representing the original example.
      <key>_position: an int32 tensor identifying the position within the original
         example.
  
      Example:
      Two input examples get combined to form an output example.
      The input examples are:
      {"inputs": [8, 7, 1, 0], "targets":[4, 1, 0]}
      {"inputs": [2, 3, 4, 1], "targets":[5, 6, 1]}
      The output example is:
      {
                     "inputs": [8, 7, 1, 2, 3, 4, 1, 0, 0, 0]
        "inputs_segmentation": [1, 1, 1, 2, 2, 2, 2, 0, 0, 0]
            "inputs_position": [0, 1, 2, 0, 1, 2, 3, 0, 0, 0]
                    "targets": [4, 1, 5, 6, 1, 0, 0, 0, 0, 0]
       "targets_segmentation": [1, 1, 2, 2, 2, 0, 0, 0, 0, 0]
           "targets_position": [0, 1, 0, 1, 2, 0, 0, 0, 0, 0]
      }
  
      0 represents padding in both the inputs and the outputs.
  
      Sequences in the incoming examples are truncated to length "length", and the
      sequences in the output examples all have fixed (padded) length "length".
  
      Args:
        dataset: a tf.data.Dataset
        length: an integer, or a dict from feature-key to integer
        keys: a collection of strings (e.g. ["inputs", "targets"])
        use_custom_ops: a boolean - custom ops are faster but require a custom-built
          binary, which is not currently possible on cloud-tpu.
  
      Returns:
        a tf.data.Dataset
      """
      shapes = tf.data.get_output_shapes(dataset)
      if keys is None:
          keys = list(shapes.keys())
      for k in keys:
          if k not in shapes:
              raise ValueError(
                  "Key %s not found in dataset.  Available keys are %s"
                  % (k, shapes.keys())
              )
          if not shapes[k].is_compatible_with(tf.TensorShape([None])):
              raise ValueError(
                  "Tensors to be packed must be one-dimensional."
              )
      # make sure that the length dictionary contains all keys as well as the
      # keys suffixed by "_segmentation" and "_position"
      length_dict = {}
      for k in keys:
          for suffix in ["", "_segmentation", "_position"]:
              length_dict[k + suffix] = (
                  length if isinstance(length, int) else length[k]
              )
      length = length_dict
  
      # trim to length
      dataset = dataset.map(
          lambda x: {k: x[k][: length[k]] for k in keys},
          num_parallel_calls=tf.data.experimental.AUTOTUNE,
      )
      # Setting batch_size=length ensures that the concatenated sequences (if they
      # have length >=1) are sufficient to fill at least one packed example.
      batch_size = max(length.values())
      dataset = dataset.padded_batch(
          batch_size, padded_shapes={k: [-1] for k in keys}
      )
      if use_custom_ops and len(keys) <= 2:
          dataset = _pack_with_custom_ops(dataset, keys, length)
      else:
          dataset = _pack_with_tf_ops(dataset, keys, length)
  
      # Set the Tensor shapes correctly since they get lost in the process.
      def my_fn(x):
          return {k: tf.reshape(v, [length[k]]) for k, v in x.items()}
  
      return dataset.map(
          my_fn, num_parallel_calls=tf.data.experimental.AUTOTUNE
      )
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import collections
  import os
  import gin
  import tensorflow.compat.v1
  import tensorflow_datasets
label_desc: Use image node_id 14 for calling a global function with example usage: pack_dataset(dataset, length, keys, use_custom_ops) and returns: dataset, unknown