code: def load_params_from_tf_checkpoint(
      checkpoint_path: str,
  ) -> tuple[
      frozen_dict.FrozenDict,
      frozen_dict.FrozenDict,
      frozen_dict.FrozenDict,
      frozen_dict.FrozenDict,
  ]:
      """Like the one with the same name plus returning the pretraining params."""
      ckpt_reader = tf.train.load_checkpoint(checkpoint_path)
      tf_params = {
          tf_name: ckpt_reader.get_tensor(tf_name)
          for tf_name in ckpt_reader.get_variable_to_dtype_map()
      }
      encoder_params = (
          bert_checkpoint_converter.convert_bert_encoder_params(
              tf_params, "bert/"
          )
      )
      pooler_params = bert_checkpoint_converter.param_conversion_util.convert_tf_params(
          tf_params,
          {("dense", "bias"): "bias", ("dense", "kernel"): "kernel"},
          "bert/pooler/dense/",
      )
      mlm_head_params = convert_bert_mlm_head_params(
          tf_params, "cls/predictions/"
      )
      nsp_head_params = bert_checkpoint_converter.param_conversion_util.convert_tf_params(
          tf_params,
          {
              (
                  "mlp",
                  "bias",
              ): "output_bias",
              (
                  "mlp",
                  "kernel",
              ): "output_weights",
          },
          "cls/seq_relationship/",
      )
      return (
          frozen_dict.freeze(encoder_params),
          frozen_dict.freeze(pooler_params),
          frozen_dict.freeze(mlm_head_params),
          frozen_dict.freeze(nsp_head_params),
      )
  
imports: import dataclasses
  from typing import Any
  from typing import Optional
  from typing import Union
  import flax
  from flax.core import frozen_dict
  import flax.linen
  import jax
  import jax.numpy
  from scenic.projects.lang4video.model.base_encoders import TextEncoder
  import tensorflow
  from flaxformer.architectures.bert import bert
  from flaxformer.architectures.bert import bert_checkpoint_converter
  import flaxformer.architectures.bert.configs
  import flaxformer.architectures.bert.heads
label_desc: Use image node_id 2 for calling a global function with example usage: load_params_from_tf_checkpoint(checkpoint_path) and returns: frozen_dict,  frozen_dict,  frozen_dict,  frozen_dict