code: def convert_pytorch_sharded_state_dict_to_flax(
      shard_filenames, flax_model
  ):
      import torch
  
      # Load the index
      flax_state_dict = {}
      for shard_file in shard_filenames:
          # load using msgpack utils
          pt_state_dict = torch.load(shard_file)
          pt_state_dict = {
              k: v.numpy() for k, v in pt_state_dict.items()
          }
  
          model_prefix = flax_model.base_model_prefix
          random_flax_state_dict = flatten_dict(flax_model.params)
  
          load_model_with_head_into_base_model = (
              model_prefix not in flax_model.params
          ) and (
              model_prefix
              in set([k.split(".")[0] for k in pt_state_dict.keys()])
          )
          load_base_model_into_model_with_head = (
              model_prefix in flax_model.params
          ) and (
              model_prefix
              not in set(
                  [k.split(".")[0] for k in pt_state_dict.keys()]
              )
          )
          # Need to change some parameters name to match Flax names
          for pt_key, pt_tensor in pt_state_dict.items():
              pt_tuple_key = tuple(pt_key.split("."))
  
              # remove base model prefix if necessary
              has_base_model_prefix = pt_tuple_key[0] == model_prefix
              if (
                  load_model_with_head_into_base_model
                  and has_base_model_prefix
              ):
                  pt_tuple_key = pt_tuple_key[1:]
  
              # Correctly rename weight parameters
              flax_key, flax_tensor = rename_key_and_reshape_tensor(
                  pt_tuple_key,
                  pt_tensor,
                  random_flax_state_dict,
                  model_prefix,
              )
              # add model prefix if necessary
              require_base_model_prefix = (
                  model_prefix,
              ) + flax_key in random_flax_state_dict
              if (
                  load_base_model_into_model_with_head
                  and require_base_model_prefix
              ):
                  flax_key = (model_prefix,) + flax_key
  
              if flax_key in random_flax_state_dict:
                  if (
                      flax_tensor.shape
                      != random_flax_state_dict[flax_key].shape
                  ):
                      raise ValueError(
                          f"PyTorch checkpoint seems to be incorrect. Weight {pt_key} was expected to be of shape "
                          f"{random_flax_state_dict[flax_key].shape}, but is {flax_tensor.shape}."
                      )
  
              # also add unexpected weight so that warning is thrown
              flax_state_dict[flax_key] = jnp.asarray(flax_tensor)
      return unflatten_dict(flax_state_dict)
  
file_path: H2O/h2o_flexgen/benchmark/third_party/transformers/src/transformers/modeling_flax_pytorch_utils.py
imports: import os
  from pickle import UnpicklingError
  from typing import Dict, Tuple
  import numpy
  import jax
  import jax.numpy
  import transformers
  from flax.serialization import from_bytes
  from flax.traverse_util import flatten_dict, unflatten_dict
  from .utils import logging
label_desc: Use image node_id 4 for calling a global function with example usage: convert_pytorch_sharded_state_dict_to_flax(shard_filenames, flax_model) and returns: unflatten_dict