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.baselines.clip import layers
  from scenic.projects.baselines.clip import download
  from tensorflow.io import gfile
label_desc: Use image node_id 12 for calling a global function with example usage: _convert_vars(torch_vars, dim_head) and returns: new_vars