code: def _pack_with_tf_ops(dataset, keys, length):
      """Helper-function for packing a dataset which has already been batched.
  
      See pack_dataset()
  
      Uses tf.while_loop.  Slow.
  
      Args:
        dataset: a dataset containing padded batches of examples.
        keys: a collection of strings
        length: an dict from feature-key to integer
  
      Returns:
        a dataset.
      """
      empty_example = {}
      for k in keys:
          empty_example[k] = tf.zeros([0], dtype=tf.int32)
          empty_example[k + "_position"] = tf.zeros([0], dtype=tf.int32)
      keys_etc = empty_example.keys()
  
      def write_packed_example(partial, outputs):
          new_partial = empty_example.copy()
          new_outputs = {}
          for k in keys_etc:
              new_outputs[k] = outputs[k].write(
                  outputs[k].size(),
                  tf.pad(
                      partial[k], [[0, length[k] - tf.size(partial[k])]]
                  ),
              )
          return new_partial, new_outputs
  
      def map_fn(x):
          """Internal function to flat_map over.
  
          Consumes a batch of input examples and produces a variable number of output
          examples.
  
          Args:
            x: a single example
          Returns:
            a tf.data.Dataset
          """
          partial = empty_example.copy()
          i = tf.zeros([], dtype=tf.int32)
          first_key, *_ = keys
          dynamic_batch_size = tf.shape(x[first_key])[0]
          outputs = {}
          for k in keys:
              outputs[k] = tf.TensorArray(
                  tf.int32,
                  size=0,
                  dynamic_size=True,
                  element_shape=[length[k]],
              )
              outputs[k + "_position"] = tf.TensorArray(
                  tf.int32,
                  size=0,
                  dynamic_size=True,
                  element_shape=[length[k]],
              )
  
          def cond_fn(i, partial, outputs):
              del partial, outputs
              return i < dynamic_batch_size
  
          def body_fn(i, partial, outputs):
              """Body function for while_loop.
  
              Args:
                i: integer scalar
                partial: dictionary of Tensor (partially-constructed example)
                outputs: dictionary of TensorArray
              Returns:
                A triple containing the new values of the inputs.
              """
              can_append = True
              one_example = {}
              for k in keys:
                  val = tf.cast(x[k][i], tf.int32)
                  val = val[
                      : tf.reduce_sum(
                          tf.cast(tf.not_equal(val, 0), tf.int32)
                      )
                  ]
                  one_example[k] = val
              for k in keys:
                  can_append = tf.logical_and(
                      can_append,
                      tf.less_equal(
                          tf.size(partial[k]) + tf.size(one_example[k]),
                          length[k],
                      ),
                  )
  
              def false_fn():
                  return write_packed_example(partial, outputs)
  
              def true_fn():
                  return partial, outputs
  
              partial, outputs = tf.cond(can_append, true_fn, false_fn)
              new_partial = {}
              for k in keys:
                  new_seq = one_example[k][: length[k]]
                  new_seq_len = tf.size(new_seq)
                  new_partial[k] = tf.concat([partial[k], new_seq], 0)
                  new_partial[k + "_position"] = tf.concat(
                      [
                          partial[k + "_position"],
                          tf.range(new_seq_len, dtype=tf.int32),
                      ],
                      0,
                  )
              partial = new_partial
              return i + 1, partial, outputs
  
          i, partial, outputs = tf.while_loop(
              cond_fn,
              body_fn,
              (i, partial, outputs),
              back_prop=False,
              shape_invariants=(
                  tf.TensorShape([]),
                  {k: tf.TensorShape([None]) for k in keys_etc},
                  {k: tf.TensorShape(None) for k in keys_etc},
              ),
          )
          partial, outputs = write_packed_example(partial, outputs)
          packed = {k: outputs[k].stack() for k in keys_etc}
          for k in keys:
              packed[k + "_segmentation"] = tf.cumsum(
                  tf.cast(
                      tf.equal(packed[k + "_position"], 0), tf.int32
                  ),
                  axis=1,
              ) * tf.cast(tf.not_equal(packed[k], 0), tf.int32)
          return packed
  
      dataset = dataset.map(
          map_fn, num_parallel_calls=tf.data.experimental.AUTOTUNE
      )
      return dataset.unbatch()
  
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 15 for calling a global function with example usage: _pack_with_tf_ops(dataset, keys, length) and returns: dataset, new_partial,  new_outputs, packed, unknown, unknown,  partial,  outputs, write_packed_example, partial,  outputs