code: def load_pytorch_state_dict_in_tf2_model(
      tf_model,
      pt_state_dict,
      tf_inputs=None,
      allow_missing_keys=False,
      output_loading_info=False,
  ):
      """Load a pytorch state_dict in a TF 2.0 model."""
      from tensorflow.python.keras import backend as K
  
      if tf_inputs is None:
          tf_inputs = tf_model.dummy_inputs
  
      if tf_inputs is not None:
          tf_model(
              tf_inputs, training=False
          )  # Make sure model is built
      # Adapt state dict - TODO remove this and update the AWS weights files instead
      # Convert old format to new format if needed from a PyTorch state_dict
      old_keys = []
      new_keys = []
      for key in pt_state_dict.keys():
          new_key = None
          if "gamma" in key:
              new_key = key.replace("gamma", "weight")
          if "beta" in key:
              new_key = key.replace("beta", "bias")
          if "running_var" in key:
              new_key = key.replace("running_var", "moving_variance")
          if "running_mean" in key:
              new_key = key.replace("running_mean", "moving_mean")
          if new_key:
              old_keys.append(key)
              new_keys.append(new_key)
      for old_key, new_key in zip(old_keys, new_keys):
          pt_state_dict[new_key] = pt_state_dict.pop(old_key)
  
      # Make sure we are able to load PyTorch base models as well as derived models (with heads)
      # TF models always have a prefix, some of PyTorch models (base ones) don't
      start_prefix_to_remove = ""
      if not any(
          s.startswith(tf_model.base_model_prefix)
          for s in pt_state_dict.keys()
      ):
          start_prefix_to_remove = tf_model.base_model_prefix + "."
  
      symbolic_weights = (
          tf_model.trainable_weights + tf_model.non_trainable_weights
      )
      tf_loaded_numel = 0
      weight_value_tuples = []
      all_pytorch_weights = set(list(pt_state_dict.keys()))
      missing_keys = []
      for symbolic_weight in symbolic_weights:
          sw_name = symbolic_weight.name
          name, transpose = convert_tf_weight_name_to_pt_weight_name(
              sw_name,
              start_prefix_to_remove=start_prefix_to_remove,
              tf_weight_shape=symbolic_weight.shape,
          )
  
          # Find associated numpy array in pytorch model state dict
          if name not in pt_state_dict:
              if allow_missing_keys:
                  missing_keys.append(name)
                  continue
              elif tf_model._keys_to_ignore_on_load_missing is not None:
                  # authorized missing keys don't have to be loaded
                  if any(
                      re.search(pat, name) is not None
                      for pat in tf_model._keys_to_ignore_on_load_missing
                  ):
                      continue
              raise AttributeError(f"{name} not found in PyTorch model")
  
          array = apply_transpose(
              transpose, pt_state_dict[name], symbolic_weight.shape
          )
  
          tf_loaded_numel += tensor_size(array)
  
          weight_value_tuples.append((symbolic_weight, array))
          all_pytorch_weights.discard(name)
  
      K.batch_set_value(weight_value_tuples)
  
      if tf_inputs is not None:
          tf_model(
              tf_inputs, training=False
          )  # Make sure restore ops are run
  
      logger.info(
          f"Loaded {tf_loaded_numel:,} parameters in the TF 2.0 model."
      )
  
      unexpected_keys = list(all_pytorch_weights)
  
      if tf_model._keys_to_ignore_on_load_missing is not None:
          for pat in tf_model._keys_to_ignore_on_load_missing:
              missing_keys = [
                  k for k in missing_keys if re.search(pat, k) is None
              ]
      if tf_model._keys_to_ignore_on_load_unexpected is not None:
          for pat in tf_model._keys_to_ignore_on_load_unexpected:
              unexpected_keys = [
                  k
                  for k in unexpected_keys
                  if re.search(pat, k) is None
              ]
  
      if len(unexpected_keys) > 0:
          logger.warning(
              "Some weights of the PyTorch model were not used when initializing the TF 2.0 model"
              f" {tf_model.__class__.__name__}: {unexpected_keys}\n- This IS expected if you are initializing"
              f" {tf_model.__class__.__name__} from a PyTorch model trained on another task or with another architecture"
              " (e.g. initializing a TFBertForSequenceClassification model from a BertForPreTraining model).\n- This IS"
              f" NOT expected if you are initializing {tf_model.__class__.__name__} from a PyTorch model that you expect"
              " to be exactly identical (e.g. initializing a TFBertForSequenceClassification model from a"
              " BertForSequenceClassification model)."
          )
      else:
          logger.warning(
              f"All PyTorch model weights were used when initializing {tf_model.__class__.__name__}.\n"
          )
      if len(missing_keys) > 0:
          logger.warning(
              f"Some weights or buffers of the TF 2.0 model {tf_model.__class__.__name__} were not initialized from the"
              f" PyTorch model and are newly initialized: {missing_keys}\nYou should probably TRAIN this model on a"
              " down-stream task to be able to use it for predictions and inference."
          )
      else:
          logger.warning(
              f"All the weights of {tf_model.__class__.__name__} were initialized from the PyTorch model.\n"
              "If your task is similar to the task the model of the checkpoint was trained on, "
              f"you can already use {tf_model.__class__.__name__} for predictions without further training."
          )
  
      if output_loading_info:
          loading_info = {
              "missing_keys": missing_keys,
              "unexpected_keys": unexpected_keys,
          }
          return tf_model, loading_info
  
      return tf_model
  
file_path: H2O/h2o_flexgen/benchmark/third_party/transformers/src/transformers/modeling_tf_pytorch_utils.py
imports: import os
  import re
  import numpy
  from .utils import ExplicitEnum, expand_dims, is_numpy_array, is_torch_tensor, logging, reshape, squeeze, tensor_size
  from .utils import transpose
label_desc: Use image node_id 6 for calling a global function with example usage: load_pytorch_state_dict_in_tf2_model(tf_model, pt_state_dict, tf_inputs, allow_missing_keys, output_loading_info) and returns: tf_model, tf_model,  loading_info