code: def initialize_bert_model(
      *,
      model_def: nn.Module,
      input_spec: Dict[
          str,
          Union[
              Tuple[Tuple[int, ...], jnp.dtype], Tuple[int, ...], None
          ],
      ],
      config: ml_collections.ConfigDict,
      rngs: Union[jnp.ndarray, Mapping[str, jnp.ndarray]],
  ) -> Tuple[PyTree, PyTree, int, Optional[float]]:
      """Initializes parameters and model state of BERT.
  
      Args:
        model_def: Definition of a model.
        input_spec: A dictionary of arg name to a (shape, dtype) pair specifying the
          shape and dtype of the input arguments. If unspecified the dtype is
          float32.
        config: Configurations of the initialization.
        rngs: Jax rng keys.
  
      Returns:
        Initial params, Init model_state, and number of trainable_params.
      """
      batch_size = (
          (config.batch_size // jax.device_count())
          if config.get("batch_size")
          else None
      )
      dummy_input = {}
      for name, spec in input_spec.items():
          if spec is not None:
              in_st = debug_utils.input_spec_to_jax_shape_dtype_struct(
                  spec, batch_size=batch_size
              )
              dummy_input[name] = jnp.zeros(in_st.shape, in_st.dtype)
          else:
              dummy_input[name] = None
  
      # We want all parameters to be created in host RAM, not on any device, they'll
      # be sent there later as needed, otherwise we already encountered two
      # situations where we allocate them twice.
      @functools.partial(jax.jit, backend="cpu")
      def _initialize_model(rngs):
          """Initialization function to be jitted."""
          init_model_state, init_params = model_def.init(
              rngs, dummy_input, train=False, debug=False
          ).pop("params")
          # Set bias in the head to low value, such that loss is small initially.
          if config.get("init_head_bias", None) is not None:
              init_params = flax.core.unfreeze(init_params)
              potential_keys = {
                  "classification_head",
                  "regression_head",
                  "next_sentence_prediction_head",
              }
              head_key = potential_keys & set(init_params.keys())
              assert len(head_key) == 1
              head_key = head_key.pop()
              new_params = optimizers.tree_map_with_names(
                  lambda p: jnp.full_like(p, config.init_head_bias),
                  init_params[head_key]["output_projection"],
                  match_name_fn=lambda name: "bias" in name,
              )
              init_params[head_key]["output_projection"] = new_params
              init_params = flax.core.freeze(init_params)
          return init_params, init_model_state
  
      if not isinstance(rngs, dict):
          rngs = {"params": rngs}
      init_params, init_model_state = _initialize_model(rngs)
      # Pop out params rng:
      rngs.pop("params")
  
      # Count number of trainable parameters:
      num_trainable_params = debug_utils.log_param_shapes(init_params)
  
      # Count gflops:
      count_flops = config.get(
          "count_flops",
          ml_collections.ConfigDict({"count_flops": True}),
      )
      if count_flops:
          variables = {"params": init_params, **init_model_state}
          flax_model_apply_fn = functools.partial(
              model_def.apply,
              variables,
              train=False,
              debug=False,
              rngs=rngs,
          )
          analysis = (
              jax.jit(flax_model_apply_fn)
              .lower(dummy_input)
              .cost_analysis()
          )
          flops = analysis["flops"]
          if count_flops.get("fuse_multiply_add", True):
              flops = flops / 2
          gflops = flops / (10**9)
          logging.info(
              "GFLOPs %0.3f for input spec: %s", gflops, input_spec
          )
      else:
          gflops = None
  
      return init_params, init_model_state, num_trainable_params, gflops
  
imports: import functools
  from typing import Any, Dict, List, Mapping, Optional, Tuple, Union
  from absl import logging
  from clu import metric_writers
  import flax
  from flax import jax_utils
  import flax.linen
  import jax
  import jax.numpy
  import jax.profiler
  import ml_collections
  import numpy
  from scenic.common_lib import debug_utils
  from scenic.train_lib import optimizers
  from scenic.train_lib import train_utils
  import scipy
  import sklearn.metrics
label_desc: Use image node_id 5 for calling a global function with example usage: initialize_bert_model() and returns: init_params,  init_model_state,  num_trainable_params,  gflops, init_params,  init_model_state