code: def load_text_params(
      params: constants.PyTree, config: ml_collections.ConfigDict
  ):
      """Load T5 params from a checkpoint."""
      t5_params = t5_model.load_pretrained_weights(config.model.t5_name)
      logging.info("T5 params are:")
      logging.info(jax.tree_util.tree_map(lambda x: x.shape, t5_params))
      # first load the shared token embeddings
  
      params = replace_dict(
          params,
          t5_params,
          ckpt_prefix_path=["params", "t5_module", "token_embedder"],
          model_prefix_path=["shared_token_embedder"],
          name_mapping={},
      )
  
      # then load the encoder and decoder weights
      params = replace_dict(
          params,
          t5_params,
          ckpt_prefix_path=["params", "t5_module", "encoder"],
          model_prefix_path=["text_encoder"],
          name_mapping={},
      )
  
      params = replace_dict(
          params,
          t5_params,
          ckpt_prefix_path=["params", "t5_module", "encoder"],
          model_prefix_path=["fusion_encoder"],
          name_mapping={},
      )
  
      params = replace_dict(
          params,
          t5_params,
          ckpt_prefix_path=["params", "t5_module", "encoder"],
          model_prefix_path=["query_head"],
          name_mapping={},
      )
  
      params = replace_dict(
          params,
          t5_params,
          ckpt_prefix_path=["params", "t5_module", "encoder"],
          model_prefix_path=["key_head"],
          name_mapping={},
      )
  
      params = replace_dict(
          params,
          t5_params,
          ckpt_prefix_path=["params", "t5_module", "decoder"],
          model_prefix_path=["out_decoder", "decoder_module"],
          name_mapping={},
      )
  
      return params
  
imports: import dataclasses
  import operator
  import re
  from typing import Any, Dict, List, Mapping, Optional, Tuple, Union
  from absl import logging
  from big_vision import utils
  from big_vision.models import vit
  from clu import metric_writers
  import flax
  from flax import jax_utils
  from flax.core import frozen_dict
  import jax
  import jax.numpy
  import ml_collections
  import optax
  from scenic.dataset_lib import dataset_utils
  from scenic.dataset_lib import datasets
  from scenic.projects.knowledge_visual_language.models import constants
  from scenic.projects.t5 import model
  from scenic.train_lib import train_utils
  from scenic.train_lib.train_utils import TrainState
label_desc: Use image node_id 9 for calling a global function with example usage: load_text_params(params, config) and returns: params