code: def add_cls_token(
      x_all_dict: ArrayDict,
      classifier: str,
      cls_encoded_dict: Optional[ArrayDict] = None,
  ) -> ArrayDict:
      """Add the cls token."""
      # If cls_encoded_dict is None, then the cls token will be generated.
      # Otherwise with be added back to the matrix.
      if cls_encoded_dict is None:
          # generate the cls token
          pass
      if classifier == "token":
          raise NotImplementedError(
              "Token classifer is not implemented yet!"
          )
  
      return x_all_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 2 for calling a global function with example usage: add_cls_token(x_all_dict, classifier, cls_encoded_dict) and returns: x_all_dict