code: def preprocess_data(
      dataset,
      shuffle: bool,
      num_epochs: Optional[int] = 1,
      pack_examples: bool = True,
      shuffle_buffer_size: int = 1024,
      max_length: int = 512,
      batch_size: int = 256,
      drop_remainder: bool = True,
      prefetch_size: int = AUTOTUNE,
  ):
      """Shuffle and batch/pack the given dataset."""
  
      def length_filter(max_len):
          def filter_fn(x):
              source, target = x["inputs"], x["targets"]
              l = tf.maximum(tf.shape(source)[0], tf.shape(target)[0])
              return tf.less(l, max_len + 1)
  
          return filter_fn
  
      if max_length > 0:
          dataset = dataset.filter(length_filter(max_length))
  
      if shuffle:
          dataset = dataset.shuffle(shuffle_buffer_size)
      dataset = dataset.repeat(num_epochs)
  
      if pack_examples:
          dataset = pack_dataset(dataset, max_length)
          dataset = dataset.batch(
              batch_size, drop_remainder=drop_remainder
          )
      else:  # simple (static-shape) padded batching
          dataset = dataset.padded_batch(
              batch_size,
              padded_shapes={
                  "inputs": max_length,
                  "targets": max_length,
              },
              padding_values={"inputs": 0, "targets": 0},
              drop_remainder=drop_remainder,
          )
  
      if prefetch_size:
          dataset = dataset.prefetch(prefetch_size)
  
      return dataset
  
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 4 for calling a global function with example usage: preprocess_data(dataset, shuffle, num_epochs, pack_examples, shuffle_buffer_size, max_length, batch_size, drop_remainder, prefetch_size) and returns: dataset, filter_fn, tf