code: def _merge_params(
      params, restored_params, model_cfg, restored_model_cfg
  ):
      """Merges `restored_params` into `params`."""
      # Start moving parameters, one-by-one and apply changes if needed.
      for m_key, m_params in restored_params.items():
          if m_key == "output_projection":
              # For the classifier head, we use a the randomly initialized params and
              #   ignore the one from pretrained model.
              pass
  
          elif m_key == "pre_logits":
              if model_cfg.model.representation_size is None:
                  # We don't have representation_size in the new model, so let's ignore
                  #   it from the pretained model, in case it has it.
                  # Note, removing the key from the dictionary is necessary to prevent
                  #   obscure errors from the Flax optimizer.
                  params.pop(m_key, None)
              else:
                  assert restored_model_cfg.model.representation_size
                  params[m_key] = m_params
  
          elif m_key == "Transformer":
              for tm_key, tm_params in m_params.items():
                  if (
                      tm_key == "posembed_input"
                  ):  # Might need resolution change.
                      posemb = params[m_key]["posembed_input"][
                          "pos_embedding"
                      ]
                      restored_posemb = m_params["posembed_input"][
                          "pos_embedding"
                      ]
  
                      if restored_posemb.shape != posemb.shape:
                          # Rescale the grid of pos, embeddings: param shape is (1, N, d).
                          logging.info(
                              "Resized variant: %s to %s",
                              restored_posemb.shape,
                              posemb.shape,
                          )
                          ntok = posemb.shape[1]
                          if (
                              restored_model_cfg.model.classifier
                              == "token"
                          ):
                              # The first token is the CLS token.
                              restored_posemb_grid = restored_posemb[
                                  0, 1:
                              ]
                              if model_cfg.model.classifier == "token":
                                  # CLS token in restored model and in target.
                                  cls_tok = restored_posemb[:, :1]
                                  ntok -= 1
                              else:
                                  # CLS token in restored model, but not target.
                                  cls_tok = restored_posemb[:, :0]
                          else:
                              restored_posemb_grid = restored_posemb[0]
                              if model_cfg.model.classifier == "token":
                                  # CLS token in target, but not restored model.
                                  cls_tok = posemb[:, :1]
                                  ntok -= 1
                              else:
                                  # CLS token not in target or restored model.
                                  cls_tok = restored_posemb[:, :0]
  
                          restored_gs = int(
                              np.sqrt(len(restored_posemb_grid))
                          )
                          gs = int(np.sqrt(ntok))
                          if (
                              restored_gs != gs
                          ):  # We need resolution change.
                              logging.info(
                                  "Grid-size from %s to %s.",
                                  restored_gs,
                                  gs,
                              )
                              restored_posemb_grid = (
                                  restored_posemb_grid.reshape(
                                      restored_gs, restored_gs, -1
                                  )
                              )
                              zoom = (
                                  gs / restored_gs,
                                  gs / restored_gs,
                                  1,
                              )
                              restored_posemb_grid = scipy.ndimage.zoom(
                                  restored_posemb_grid, zoom, order=1
                              )
                          # Attach the CLS token again.
                          restored_posemb_grid = (
                              restored_posemb_grid.reshape(
                                  1, gs * gs, -1
                              )
                          )
                          restored_posemb = jnp.array(
                              np.concatenate(
                                  [cls_tok, restored_posemb_grid],
                                  axis=1,
                              )
                          )
  
                      params[m_key][tm_key][
                          "pos_embedding"
                      ] = restored_posemb
                  # Other parameters of the Transformer encoder if they are in the target.
                  elif tm_key in params[m_key]:
                      params[m_key][tm_key] = tm_params
                  else:
                      logging.info(
                          "Ignoring %s. In restored model's Transformer,"
                          "but not in target",
                          m_key,
                      )
  
          elif m_key in params:
              # Use the rest if they are in the pretrained model.
              params[m_key] = m_params
  
          else:
              logging.info(
                  "Ignoring %s. In restored model, but not in target",
                  m_key,
              )
  
imports: from typing import Any, Callable, Optional, Sequence
  from absl import logging
  import flax
  import flax.linen
  import jax
  import jax.numpy
  import ml_collections
  import numpy
  from scenic.model_lib.base_models.multilabel_classification_model import MultiLabelClassificationModel
  from scenic.model_lib.layers import attention_layers
  from scenic.model_lib.layers import nn_layers
  import scipy
  from tensorflow.io import gfile
label_desc: Use image node_id 2 for calling a global function with example usage: _merge_params(params, restored_params, model_cfg, restored_model_cfg) without return types