code: def adapt_attention_weights(
      params: Dict[str, Array], encoder_num_heads: int
  ) -> Array:
      """Adapts PyTorch attention layers to Flax."""
  
      new_params = {}
      for name, parameter in params.items():
          if "qkv.weight" in name or "proj.weight" in name:
              if "encoder" in name:
                  num_heads = encoder_num_heads
              elif "decoder" in name:
                  continue
              else:
                  raise ValueError(f"Unsupported parameter {name}.")
  
          if "attn.qkv.weight" in name:
              # Un-merge the query-, key- and value-projections.
              query, key, value = np.split(parameter, 3, axis=0)
              query, key, value = (
                  query.transpose(),
                  key.transpose(),
                  value.transpose(),
              )
              d, _ = query.shape
              query = query.reshape(d, num_heads, -1)
              key = key.reshape(d, num_heads, -1)
              value = value.reshape(d, num_heads, -1)
  
              prefix = name.replace("qkv.weight", "")
              new_params[prefix + "q_weight"] = query
              new_params[prefix + "k_weight"] = key
              new_params[prefix + "v_weight"] = value
  
              # We also need to reshape the bias parameters.
              q_bias = params[prefix + "q_bias"]
              q_bias = np.reshape(q_bias, (num_heads, -1))
  
              v_bias = params[prefix + "v_bias"]
              v_bias = np.reshape(v_bias, (num_heads, -1))
  
              if prefix + "k_bias" in params:
                  k_bias = params[prefix + "v_bias"]
                  k_bias = np.reshape(k_bias, (num_heads, -1))
              else:
                  logging.log_first_n(
                      logging.INFO,
                      "No key bias. Initialising with zeros",
                      1,
                  )
                  k_bias = np.zeros(q_bias.shape)
  
              new_params[prefix + "q_bias"] = q_bias
              new_params[prefix + "v_bias"] = v_bias
              new_params[prefix + "k_bias"] = k_bias
  
          elif "attn.proj.weight" in name:
              # Reshape the output projection with the number of heads.
              d, _ = parameter.shape
              output_proj = parameter.transpose().reshape(
                  (num_heads, -1, d)
              )
              new_params[name] = output_proj
  
          else:
              new_params[name] = parameter
  
      return new_params
  
imports: import pickle
  from typing import Any, Dict, Union
  from absl import app
  from absl import flags
  from absl import logging
  import flax
  import jax
  from jax import random
  import jax.numpy
  import ml_collections
  import numpy
  from scenic.projects.objectvivit import model
  from scenic.train_lib import train_utils
label_desc: Use image node_id 3 for calling a global function with example usage: adapt_attention_weights(params, encoder_num_heads) and returns: new_params