code: def initialize_model(
      *,
      model_def: nn.Module,
      input_spec: Mapping[str, Tuple[Tuple[int, ...], jnp.dtype]],
      config: ml_collections.ConfigDict,
      rngs: Union[jnp.ndarray, Mapping[str, jnp.ndarray]],
  ) -> Tuple[PyTree, PyTree, int, Optional[float]]:
      """Initializes parameters and model state.
  
      This function is branched from scenic/train_lib_deprecated/train_utils.py.
      The difference is that here the input_spec is a dictionary while it is a
      sequence in the original implementation.
  
      Args:
        model_def: Definition of a model.
        input_spec: An mapping of modality name to (shape, dtype) pairs specifying
          the shape and dtype of the inputs.
        config: Configurations of the initialization.
        rngs: Jax rng keys.
  
      Returns:
        Initial params, initial model_state, and number of trainable_params.
      """
      batch_size = (
          (config.batch_size // jax.device_count())
          if config.get("batch_size")
          else None
      )
      input_placeholder = {}
      for modality_name, spec in input_spec.items():
          in_st = debug_utils.input_spec_to_jax_shape_dtype_struct(
              spec, batch_size=batch_size
          )
          input_placeholder[modality_name] = jnp.zeros(
              in_st.shape, in_st.dtype
          )
  
      # 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 = flax.core.pop(
              model_def.init(
                  rngs, input_placeholder, train=False, debug=False
              ),
              "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)
              init_params[
                  "output_projection"
              ] = optimizers.tree_map_with_names(
                  lambda p: jnp.full_like(p, config.init_head_bias),
                  init_params["output_projection"],
                  match_name_fn=lambda name: "bias" in name,
              )
              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}
          flops = compute_flops(
              flax_model_apply_fn=functools.partial(
                  model_def.apply,
                  variables,
                  train=False,
                  debug=False,
                  rngs=rngs,
              ),
              input_spec=count_flops.get("input_spec", input_spec),
              fuse_multiply_add=count_flops.get(
                  "fuse_multiply_add", True
              ),
          )
          gflops = flops / (10**9)
      else:
          gflops = None
  
      return init_params, init_model_state, num_trainable_params, gflops
  
file_path: scenic/scenic/projects/mbt/train_utils.py
imports: import functools
  from typing import Any, Callable, Tuple, Optional, Mapping, Union
  from absl import logging
  import flax
  import flax.linen
  import jax
  import jax.numpy
  import ml_collections
  from scenic.common_lib import debug_utils
  from scenic.train_lib_deprecated import optimizers
label_desc: Use image node_id 2 for calling a global function with example usage: initialize_model() and returns: init_params,  init_model_state,  num_trainable_params,  gflops, init_params,  init_model_state