code: def convert_pytorch_state_dict_to_flax(
      pt_state_dict, flax_model, init_key=42
  ):
      # Step 1: Convert pytorch tensor to numpy
      pt_state_dict = {k: v.numpy() for k, v in pt_state_dict.items()}
  
      # Step 2: Since the model is stateless, get random Flax params
      random_flax_params = flax_model.init_weights(PRNGKey(init_key))
  
      random_flax_state_dict = flatten_dict(random_flax_params)
      flax_state_dict = {}
  
      # Need to change some parameters name to match Flax names
      for pt_key, pt_tensor in pt_state_dict.items():
          renamed_pt_key = rename_key(pt_key)
          pt_tuple_key = tuple(renamed_pt_key.split("."))
  
          # Correctly rename weight parameters
          flax_key, flax_tensor = rename_key_and_reshape_tensor(
              pt_tuple_key, pt_tensor, random_flax_state_dict
          )
  
          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: llmga/llmga/diffusers/src/diffusers/models/modeling_flax_pytorch_utils.py
imports: import re
  import jax.numpy
  from flax.traverse_util import flatten_dict, unflatten_dict
  from jax.random import PRNGKey
  from ..utils import logging
label_desc: Use image node_id 3 for calling a global function with example usage: convert_pytorch_state_dict_to_flax(pt_state_dict, flax_model, init_key) and returns: unflatten_dict