code: def _load_experts(config, prefix, mat, weights):
      if config.quantize is not None:
          raise NotImplementedError(
              "Mixtral does not support weight quantization yet."
          )
  
      assert mat in ["w1", "w2", "w3"]
  
      world_size = weights.process_group.size()
      rank = weights.process_group.rank()
  
      assert (
          config.intermediate_size % world_size == 0
      ), f"The chosen size {config.intermediate_size} is not compatible with sharding on {world_size} shards"
  
      block_size = config.intermediate_size // world_size
      start = rank * block_size
      stop = (rank + 1) * block_size
  
      tensor = torch.empty(
          (config.num_local_experts * block_size, config.hidden_size),
          dtype=weights.dtype,
          device=weights.device,
      )
  
      for i in range(config.num_local_experts):
          slice_ = weights._get_slice(f"{prefix}.{i}.{mat}.weight")
  
          if mat == "w2":
              expert_slice = slice_[:, start:stop].t().contiguous()
          else:
              expert_slice = slice_[start:stop]
          tensor[
              i * block_size : (i + 1) * block_size
          ] = expert_slice.to(dtype=weights.dtype).to(
              device=weights.device
          )
      return tensor
  
imports: import torch
  import torch.distributed
  import numpy
  from torch import nn
  from transformers.activations import ACT2FN
  from transformers.configuration_utils import PretrainedConfig
  from typing import Optional, List, Tuple
  from loguru import logger
  from text_generation_server.utils import paged_attention, flash_attn
  from text_generation_server.utils.layers import FastLinear, FastRMSNorm, TensorParallelRowLinear, TensorParallelColumnLinear, TensorParallelEmbedding, PositionRotaryEmbedding, TensorParallelHead, get_linear
label_desc: Use image node_id 4 for calling a global function with example usage: _load_experts(config, prefix, mat, weights) and returns: tensor