code: def _load_weights(
      model: VisionTransformer, checkpoint_path: str, prefix: str = ""
  ):
      """Load weights from .npz checkpoints for official Google Brain Flax implementation"""
      import numpy as np
  
      def _n2p(w, t=True):
          if (
              w.ndim == 4
              and w.shape[0] == w.shape[1] == w.shape[2] == 1
          ):
              w = w.flatten()
          if t:
              if w.ndim == 4:
                  w = w.transpose([3, 2, 0, 1])
              elif w.ndim == 3:
                  w = w.transpose([2, 0, 1])
              elif w.ndim == 2:
                  w = w.transpose([1, 0])
          return torch.from_numpy(w)
  
      w = np.load(checkpoint_path)
      if not prefix and "opt/target/embedding/kernel" in w:
          prefix = "opt/target/"
  
      if hasattr(model.patch_embed, "backbone"):
          # hybrid
          backbone = model.patch_embed.backbone
          stem_only = not hasattr(backbone, "stem")
          stem = backbone if stem_only else backbone.stem
          stem.conv.weight.copy_(
              adapt_input_conv(
                  stem.conv.weight.shape[1],
                  _n2p(w[f"{prefix}conv_root/kernel"]),
              )
          )
          stem.norm.weight.copy_(_n2p(w[f"{prefix}gn_root/scale"]))
          stem.norm.bias.copy_(_n2p(w[f"{prefix}gn_root/bias"]))
          if not stem_only:
              for i, stage in enumerate(backbone.stages):
                  for j, block in enumerate(stage.blocks):
                      bp = f"{prefix}block{i + 1}/unit{j + 1}/"
                      for r in range(3):
                          getattr(block, f"conv{r + 1}").weight.copy_(
                              _n2p(w[f"{bp}conv{r + 1}/kernel"])
                          )
                          getattr(block, f"norm{r + 1}").weight.copy_(
                              _n2p(w[f"{bp}gn{r + 1}/scale"])
                          )
                          getattr(block, f"norm{r + 1}").bias.copy_(
                              _n2p(w[f"{bp}gn{r + 1}/bias"])
                          )
                      if block.downsample is not None:
                          block.downsample.conv.weight.copy_(
                              _n2p(w[f"{bp}conv_proj/kernel"])
                          )
                          block.downsample.norm.weight.copy_(
                              _n2p(w[f"{bp}gn_proj/scale"])
                          )
                          block.downsample.norm.bias.copy_(
                              _n2p(w[f"{bp}gn_proj/bias"])
                          )
          embed_conv_w = _n2p(w[f"{prefix}embedding/kernel"])
      else:
          embed_conv_w = adapt_input_conv(
              model.patch_embed.proj.weight.shape[1],
              _n2p(w[f"{prefix}embedding/kernel"]),
          )
      model.patch_embed.proj.weight.copy_(embed_conv_w)
      model.patch_embed.proj.bias.copy_(
          _n2p(w[f"{prefix}embedding/bias"])
      )
      model.cls_token.copy_(_n2p(w[f"{prefix}cls"], t=False))
      pos_embed_w = _n2p(
          w[f"{prefix}Transformer/posembed_input/pos_embedding"],
          t=False,
      )
      if pos_embed_w.shape != model.pos_embed.shape:
          pos_embed_w = resize_pos_embed(  # resize pos embedding when different size from pretrained weights
              pos_embed_w,
              model.pos_embed,
              getattr(model, "num_tokens", 1),
              model.patch_embed.grid_size,
          )
      model.pos_embed.copy_(pos_embed_w)
      model.norm.weight.copy_(
          _n2p(w[f"{prefix}Transformer/encoder_norm/scale"])
      )
      model.norm.bias.copy_(
          _n2p(w[f"{prefix}Transformer/encoder_norm/bias"])
      )
      if (
          isinstance(model.head, nn.Linear)
          and model.head.bias.shape[0]
          == w[f"{prefix}head/bias"].shape[-1]
      ):
          model.head.weight.copy_(_n2p(w[f"{prefix}head/kernel"]))
          model.head.bias.copy_(_n2p(w[f"{prefix}head/bias"]))
      if (
          isinstance(getattr(model.pre_logits, "fc", None), nn.Linear)
          and f"{prefix}pre_logits/bias" in w
      ):
          model.pre_logits.fc.weight.copy_(
              _n2p(w[f"{prefix}pre_logits/kernel"])
          )
          model.pre_logits.fc.bias.copy_(
              _n2p(w[f"{prefix}pre_logits/bias"])
          )
      for i, block in enumerate(model.blocks.children()):
          block_prefix = f"{prefix}Transformer/encoderblock_{i}/"
          mha_prefix = block_prefix + "MultiHeadDotProductAttention_1/"
          block.norm1.weight.copy_(
              _n2p(w[f"{block_prefix}LayerNorm_0/scale"])
          )
          block.norm1.bias.copy_(
              _n2p(w[f"{block_prefix}LayerNorm_0/bias"])
          )
          block.attn.qkv.weight.copy_(
              torch.cat(
                  [
                      _n2p(w[f"{mha_prefix}{n}/kernel"], t=False)
                      .flatten(1)
                      .T
                      for n in ("query", "key", "value")
                  ]
              )
          )
          block.attn.qkv.bias.copy_(
              torch.cat(
                  [
                      _n2p(w[f"{mha_prefix}{n}/bias"], t=False).reshape(
                          -1
                      )
                      for n in ("query", "key", "value")
                  ]
              )
          )
          block.attn.proj.weight.copy_(
              _n2p(w[f"{mha_prefix}out/kernel"]).flatten(1)
          )
          block.attn.proj.bias.copy_(_n2p(w[f"{mha_prefix}out/bias"]))
          for r in range(2):
              if hasattr(block.mlp, "experts"):
                  num_experts = getattr(block.mlp, "experts").w1.shape[
                      0
                  ]
                  for ind in range(num_experts):
                      getattr(block.mlp.experts, f"w{r + 1}")[
                          ind
                      ].copy_(
                          torch.from_numpy(
                              w[
                                  f"{block_prefix}MlpBlock_3/Dense_{r}/kernel"
                              ]
                          )
                      )
              else:
                  getattr(block.mlp, f"fc{r + 1}").weight.copy_(
                      _n2p(
                          w[
                              f"{block_prefix}MlpBlock_3/Dense_{r}/kernel"
                          ]
                      )
                  )
                  getattr(block.mlp, f"fc{r + 1}").bias.copy_(
                      _n2p(
                          w[f"{block_prefix}MlpBlock_3/Dense_{r}/bias"]
                      )
                  )
          block.norm2.weight.copy_(
              _n2p(w[f"{block_prefix}LayerNorm_2/scale"])
          )
          block.norm2.bias.copy_(
              _n2p(w[f"{block_prefix}LayerNorm_2/bias"])
          )
  
imports: import logging
  import os
  import math
  from collections import OrderedDict
  from copy import deepcopy
  from typing import Any, Callable, Optional, Tuple
  import torch
  import torch.nn
  from torch.hub import load_state_dict_from_url
  from timm.models.features import FeatureListNet, FeatureDictNet, FeatureHookNet
  from timm.models.fx_features import FeatureGraphNet
  from timm.models.hub import has_hf_hub, download_cached_file, load_state_dict_from_hf
  from timm.models.layers import Conv2dSame, Linear
  from timm.models.vision_transformer import VisionTransformer, resize_pos_embed
  import tutel.impls.moe_layer
label_desc: Use image node_id 6 for calling a global function with example usage: _load_weights(model, checkpoint_path, prefix) and returns: torch