<SYSTEM_TASK:>
Solve the following problem using Python, implementing the functions described below, one line at a time
<END_TASK>
<USER_TASK:>
Description:
 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 keys: a list 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 = dataset.output_shapes
  if keys is None:
    keys = 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.")

  # trim to length
  dataset = dataset.map(lambda x: {k: x[k][:length] 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 = length
  dataset = dataset.padded_batch(
      batch_size, padded_shapes={k: [-1] for k in keys})
  if use_custom_ops and len(keys) <= 2:
    return _pack_with_custom_ops(dataset, keys, length)
  else:
    return _pack_with_tf_ops(dataset, keys, length)