code: def tokenize_input(
      inputs: ArrayDict,
      temporal_encoding_config: str,
      patches: ml_collections.ConfigDict,
      hidden_size: int,
  ) -> ArrayDict:
      """Tokenize the input based on their modality."""
      embed_2d = {"spectrogram"}
      temporal_encode = {
          "RGB",
          "modis",
          "l7",
          "s2",
          "s1",
          "nicfi",
          "nicfi_monthly",
          "alos",
      }
      if (
          "size" in patches
      ):  # Handles case of one patch given for all modalities.
          modal_patches = {key: patches for key in inputs}
      else:
          modal_patches = patches
      tokens_dict = {}
      for key in inputs:
          # Shape is [batch, num_tokens, hidden_size].
          if key in temporal_encode:
              x_tokens, _ = vivit_model.temporal_encode(
                  inputs[key],
                  temporal_encoding_config,
                  patches=modal_patches[key],
                  hidden_size=hidden_size,
                  name=f"embedding_{key}",
              )
          elif key in embed_2d:
              x_tokens = model_utils.embed_2d_patch(
                  inputs[key],
                  patches=modal_patches[key],
                  embedding_dim=hidden_size,
                  name=f"embedding_{key}",
              )
          else:
              raise ValueError(f"Modality {key} is not supported!")
          tokens_dict[key] = x_tokens
  
      return tokens_dict
  
imports: import functools
  from typing import Any, Dict, Optional, Tuple
  import flax.linen
  import jax
  import jax.numpy
  import ml_collections
  from scenic.model_lib.base_models import base_model
  from scenic.model_lib.base_models import model_utils
  from scenic.model_lib.layers import nn_layers
  from scenic.projects.av_mae import base_model
  from scenic.projects.av_mae import model_utils
  from scenic.projects.mbt import model
  from scenic.projects.vivit import model
  from scenic.projects.vivit import model_utils
label_desc: Use image node_id 4 for calling a global function with example usage: tokenize_input(inputs, temporal_encoding_config, patches, hidden_size) and returns: tokens_dict