code: def _pack_with_tf_ops(
      dataset: tf.data.Dataset,
      keys: List[str],
      key2length: Dict[str, int],
  ) -> tf.data.Dataset:
      """Helper-function for packing a dataset which has already been batched.
  
      Helper for pack_dataset()  Uses tf.while_loop.
  
      Args:
        dataset: a dataset containing padded batches of examples.
        keys: a list of strings
        key2length: a 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, key2length[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)
          dynamic_batch_size = tf.shape(x[keys[0]])[0]
          outputs = {}
          for k in keys:
              outputs[k] = tf.TensorArray(
                  tf.int32,
                  size=0,
                  dynamic_size=True,
                  element_shape=[key2length[k]],
              )
              outputs[k + "_position"] = tf.TensorArray(
                  tf.int32,
                  size=0,
                  dynamic_size=True,
                  element_shape=[key2length[k]],
              )
  
          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]),
                          key2length[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][: key2length[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)],
                      0,
                  )
              partial = new_partial
              return i + 1, partial, outputs
  
          # For loop over all examples in the batch.
          i, partial, outputs = tf.while_loop(
              cond=lambda *_: True,
              body=body_fn,
              loop_vars=(i, partial, outputs),
              shape_invariants=(
                  tf.TensorShape([]),
                  {k: tf.TensorShape([None]) for k in keys_etc},  # type: ignore[wrong-arg-types]
                  {k: tf.TensorShape(None) for k in keys_etc},  # type: ignore[wrong-arg-types]
              ),
              maximum_iterations=dynamic_batch_size,
          )
          _, 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=AUTOTUNE)
      return dataset.unbatch()
  
file_path: flax/flax/experimental/nnx/examples/lm1b/input_pipeline.py
imports: import os
  from typing import Dict, List, Optional, Union
  import tensorflow
  import tensorflow_datasets
  import tokenizer
  from clu import deterministic_data
  from configs import default
label_desc: Use image node_id 3 for calling a global function with example usage: _pack_with_tf_ops(dataset, keys, key2length) and returns: dataset, new_partial,  new_outputs, packed, unknown,  partial,  outputs, write_packed_example, partial,  outputs