code: def _convert_vars(
      torch_vars: Mapping[str, np.ndarray], dim_head: int = 64
  ) -> PyTree:
      """Convert torch parameters to flax parameters."""
      # Expand QKV dense input projection to separate Q, K, V projections
      # and fix shape/transposing of attention layers.
      torch_vars = _convert_attn_layers(torch_vars, dim_head)
      flax_vars = {}
      torch_vars.pop("context_length", None)
      torch_vars.pop("input_resolution", None)
      torch_vars.pop("vocab_size", None)
      for torch_key, v in torch_vars.items():
          if "num_batches_tracked" in torch_key:
              continue
  
          if "conv" in torch_key or "downsample.0.weight" in torch_key:
              v = v.transpose(2, 3, 1, 0)
          elif (
              "weight" in torch_key
              and v.ndim == 2
              and "embedding" not in torch_key
          ):
              # Fully connected layers are transposed, embeddings are not
              v = v.T
  
          jax_key = torch_key.replace(
              "visual.proj", "visual.proj.kernel"
          )
          jax_key = jax_key.replace(
              "text_projection", "text_projection.kernel"
          )
          if (
              "bn" in jax_key
              or "ln" in jax_key
              or "downsample.1" in jax_key
          ):
              jax_key = jax_key.replace(".weight", ".scale")
          else:
              jax_key = jax_key.replace(".weight", ".kernel")
          if (
              jax_key.startswith("transformer")
              or jax_key.startswith("text_projection")
              or jax_key.startswith("ln_final")
              or jax_key.startswith("positional_embedding")
          ):
              jax_key = "text." + jax_key
  
          jax_key = jax_key.replace(
              "token_embedding.kernel", "text.token_embedding.embedding"
          )
  
          jax_key = jax_key.replace(
              "attnpool.k_proj", "attnpool.attn.key"
          )
          jax_key = jax_key.replace(
              "attnpool.q_proj", "attnpool.attn.query"
          )
          jax_key = jax_key.replace(
              "attnpool.v_proj", "attnpool.attn.value"
          )
          jax_key = jax_key.replace(
              "attnpool.c_proj", "attnpool.attn.out"
          )
          if "attnpool.attn.out" in jax_key:
              if jax_key.endswith("kernel"):
                  v = v.reshape(-1, dim_head, v.shape[-1])
          elif "attnpool.attn" in jax_key:
              if jax_key.endswith("bias"):
                  v = v.reshape(-1, dim_head)
              else:
                  v = v.reshape(v.shape[0], -1, dim_head)
  
          if jax_key.endswith("running_mean"):
              jax_key = "batch_stats." + jax_key.replace(
                  ".running_mean", ".mean"
              )
          elif jax_key.endswith("running_var"):
              jax_key = "batch_stats." + jax_key.replace(
                  ".running_var", ".var"
              )
          else:
              jax_key = "params." + jax_key
  
          jax_key = jax_key.replace(".", "/")
          jax_key = jax_key.replace("resblocks/", "resblocks.")
          jax_key = jax_key.replace("resblocks/", "resblocks.")
  
          flax_vars[tuple(jax_key.split("/"))] = jnp.asarray(v)
  
      # Transform the flattened param dict to the original nested structure.
      new_vars = flax.core.freeze(
          flax.traverse_util.unflatten_dict(flax_vars)
      )
      return new_vars
  
imports: import os
  from typing import Any, Mapping, Optional
  from absl import logging
  import flax
  import jax
  import jax.numpy
  import numpy
  from scenic.projects.owl_vit.clip import layers
  from scenic.projects.baselines.clip import download
  from tensorflow.io import gfile
label_desc: Use image node_id 11 for calling a global function with example usage: _convert_vars(torch_vars, dim_head) and returns: new_vars