code: def map_generation_split(
      batch,
      span_len,
      output_max_len,
      split_key="tokens",
      add_retr=False,
  ):
      """Split tokens into prefix, decoder_input and decoder_output."""
      full_tokens = batch.pop(split_key)
      full_masks = tf.greater(full_tokens, 0)
      min_length = tf.reduce_max(
          [
              tf.reduce_min(
                  tf.reduce_sum(tf.cast(full_masks, tf.int16), axis=1)
              ),
              PROMPT_LENGTH + 4,
          ]
      ).numpy()
      max_length = PROMPT_LENGTH + span_len + 1
      bsz = full_tokens.shape[0]
      idx = tf.experimental.numpy.random.randint(
          low=PROMPT_LENGTH,
          high=tf.reduce_min([min_length, max_length]).numpy(),
      ).numpy()
      input_tokens = [
          full_tokens[..., :idx],
          tf.ones([bsz, 1], dtype=tf.int16) * MASK_TOKEN_ID,
          tf.zeros([bsz, max_length - idx - 1], dtype=tf.int16),
      ]
      output_tokens = [
          tf.ones([bsz, 1], dtype=tf.int16) * BOS_ID,
          full_tokens[..., idx : idx + output_max_len],
      ]
      batch["encoder_input_tokens"] = tf.concat(input_tokens, axis=1)
      batch["encoder_input_image"] = batch.pop("image")
      output_tokens = tf.concat(output_tokens, axis=1)
      batch["decoder_input_tokens"] = output_tokens[..., :-1]
      batch["decoder_target_tokens"] = output_tokens[..., 1:]
      if add_retr:
          if "retr_texts" in batch:
              batch["retr_texts"] = tf.expand_dims(
                  batch["retr_texts"], axis=1
              )
          else:
              batch["retr_texts"] = tf.expand_dims(
                  batch["decoder_input_tokens"], axis=1
              )
      return batch
  
imports: import functools
  from big_vision.datasets.imagenet import class_names
  import numpy
  from scenic.dataset_lib import dataset_utils
  from scenic.dataset_lib.big_transfer import registry
  from scenic.dataset_lib.big_transfer.preprocessing import utils
  from scenic.projects.t5 import tokenizer
  import tensorflow
label_desc: Use image node_id 8 for calling a global function with example usage: map_generation_split(batch, span_len, output_max_len, split_key, add_retr) and returns: batch