code: def init_bert_from_train_state(
      train_state: Any,
      restored_train_state: Any,
      config: ml_collections.ConfigDict,
      restored_model_cfg: ml_collections.ConfigDict,
  ) -> Any:
      """Updates the train_state with data from restored_train_state."""
      del restored_model_cfg
  
      def _get_param_dict(params):
          return {
              "/".join([str(kk) for kk in k]): v
              for k, v in flax.traverse_util.flatten_dict(
                  params
              ).items()
          }
  
      if hasattr(train_state, "optimizer"):
          # TODO(dehghani): Remove support for flax optim.
          params = flax.core.unfreeze(train_state.optimizer.target)
          restored_params = flax.core.unfreeze(
              restored_train_state.optimizer.target
          )
      else:
          params = flax.core.unfreeze(train_state.params)
          restored_params = flax.core.unfreeze(
              restored_train_state.params
          )
  
      params_dict = _get_param_dict(params)
      # Fix some names:
      restored_params_dict = dict()
      for key, value in flax.traverse_util.flatten_dict(
          restored_params
      ).items():
          name = "/".join([str(k) for k in key])
          if (
              config.init_from.restore_next_sentence_prediction_head_params
          ):
              name = name.replace(
                  "next_sentence_prediction_head", "classification_head"
              )
          restored_params_dict[name] = value
      # Copy parameters over:
      for pname, pvalue in restored_params_dict.items():
          if "masked_language_model_head" in pname:
              # We throw away parameters of `masked_language_model_head`, but
              # for the `next_sentence_prediction_head`, we only discard the final
              # dense (`output_projection`) tha maps model representation to the
              # label space.
              continue
          if (
              not config.init_from.restore_next_sentence_prediction_head_params
              and "next_sentence_prediction_head" in pname
          ):
              continue
          if "output_projection" in pname:
              continue
          elif pname in params_dict:
              params_dict[pname] = pvalue
          else:
              logging.error(
                  "Restored key doesn't exist in the model: %s.", pname
              )
  
      logging.info(
          "Inspect missing keys from the restored params:\n%s",
          params_dict.keys() - restored_params_dict.keys(),
      )
      logging.info(
          "Inspect extra keys the the restored params:\n%s",
          restored_params_dict.keys() - params_dict.keys(),
      )
  
      splitkeys = {
          tuple(k.split("/")): v for k, v in params_dict.items()
      }
      params = flax.traverse_util.unflatten_dict(splitkeys)
      logging.info(
          "Parameter summary after initialising from train state:"
      )
      debug_utils.log_param_shapes(params)
      if hasattr(train_state, "optimizer"):
          # TODO(dehghani): Remove support for flax optim.
          return train_state.replace(
              optimizer=train_state.optimizer.replace(
                  target=flax.core.freeze(params)
              )
          )
      else:
          return train_state.replace(params=flax.core.freeze(params))
  
imports: from typing import Any, Dict
  from absl import logging
  import flax
  import flax.linen
  import jax.numpy
  import ml_collections
  from scenic.common_lib import debug_utils
  from scenic.model_lib.base_models.classification_model import ClassificationModel
  from scenic.model_lib.base_models.regression_model import RegressionModel
  from scenic.projects.baselines.bert import bert_base_model
  from scenic.projects.baselines.bert import layers
label_desc: Use image node_id 1 for calling a global function with example usage: init_bert_from_train_state(train_state, restored_train_state, config, restored_model_cfg) and returns: unknown, train_state, train_state