class_docstr: A torch module that performs Mixture-of-Experts inference with a local gating function and multiple remote experts.
  Natively supports pytorch autograd.
  
  :note: By default, not all experts are guaranteed to perform forward pass. Moreover, not all of those who ran
   forward pass are guaranteed to perform backward pass. In the latter case, gradient will be averaged without
   the missing experts
  
  :param in_features: common input size for experts and gating function
  :param grid_size: dimensions that form expert uid (see below)
  :param uid_prefix: common prefix for all expert uids (must end with '.')
  :note: expert uid follows the pattern {uid_prefix}.{0...grid_size[0]}.{0...grid_size[1]}...{0...grid_size[-1]}
  :param dht: a DHT instance used to search for best experts
  :param k_best: average this many highest-scoring experts to compute activations
  :param k_min: make sure at least this many experts returned output (i.e. didn't fail)
  :param timeout_after_k_min: wait for this many seconds after k_min experts returned results.
   Any expert that didn't manage to return output after that delay is considered unavailable
  :param detect_anomalies: whether to check input/output tensors for NaN and infinity values
  :param allow_zero_outputs: whether to return zeros if no experts respond on forward pass
class_docstr_tok: ["A","torch","module","that","performs","Mixture-of-Experts","inference","with","a","local","gating","function","and","multiple","remote","experts",".","Natively","supports","pytorch","autograd",".",":","note",":","By","default",",","not","all","experts","are","guaranteed","to","perform","forward","pass",".","Moreover",",","not","all","of","those","who","ran","forward","pass","are","guaranteed","to","perform","backward","pass",".","In","the","latter","case",",","gradient","will","be","averaged","without","the","missing","experts",":","param","in_features",":","common","input","size","for","experts","and","gating","function",":","param","grid_size",":","dimensions","that","form","expert","uid","(","see","below",")",":","param","uid_prefix",":","common","prefix","for","all","expert","uids","(","must","end","with","'",".","'",")",":","note",":","expert","uid","follows","the","pattern","{","uid_prefix","}",".","{","0","...","grid_size","[","0","]","}",".","{","0","...","grid_size","[","1","]","}","...","{","0","...","grid_size","[","-1","]","}",":","param","dht",":","a","DHT","instance","used","to","search","for","best","experts",":","param","k_best",":","average","this","many","highest-scoring","experts","to","compute","activations",":","param","k_min",":","make","sure","at","least","this","many","experts","returned","output","(","i.e",".","did","n't","fail",")",":","param","timeout_after_k_min",":","wait","for","this","many","seconds","after","k_min","experts","returned","results",".","Any","expert","that","did","n't","manage","to","return","output","after","that","delay","is","considered","unavailable",":","param","detect_anomalies",":","whether","to","check","input\/output","tensors","for","NaN","and","infinity","values",":","param","allow_zero_outputs",":","whether","to","return","zeros","if","no","experts","respond","on","forward","pass"]
docstr: Choose k best experts with beam search, then call chosen experts and average their outputs.
  Input tensor is averaged over all dimensions except for first and last
  (we assume that extra dimensions represent sequence length or image height/width)
  
  :param input: a tensor of values that are used to estimate gating function, batch-first.
  :param args: extra positional parameters that will be passed to each expert after input, batch-first
  :param kwargs: extra keyword parameters that will be passed to each expert, batch-first
  :returns: averaged predictions of all experts that delivered result on time, nested structure of batch-first
docstr_tok: ["Choose","k","best","experts","with","beam","search",",","then","call","chosen","experts","and","average","their","outputs",".","Input","tensor","is","averaged","over","all","dimensions","except","for","first","and","last","(","we","assume","that","extra","dimensions","represent","sequence","length","or","image","height\/width",")",":","param","input",":","a","tensor","of","values","that","are","used","to","estimate","gating","function",",","batch-first",".",":","param","args",":","extra","positional","parameters","that","will","be","passed","to","each","expert","after","input",",","batch-first",":","param","kwargs",":","extra","keyword","parameters","that","will","be","passed","to","each","expert",",","batch-first",":","returns",":","averaged","predictions","of","all","experts","that","delivered","result","on","time",",","nested","structure","of","batch-first"]
code: def forward(
      self,
      input: torch.Tensor,
      *args: torch.Tensor,
      **kwargs: torch.Tensor,
  ):
      """
      Choose k best experts with beam search, then call chosen experts and average their outputs.
      Input tensor is averaged over all dimensions except for first and last
      (we assume that extra dimensions represent sequence length or image height/width)
  
      :param input: a tensor of values that are used to estimate gating function, batch-first.
      :param args: extra positional parameters that will be passed to each expert after input, batch-first
      :param kwargs: extra keyword parameters that will be passed to each expert, batch-first
      :returns: averaged predictions of all experts that delivered result on time, nested structure of batch-first
      """
      if input.ndim != 2:
          input_for_gating = input.mean(
              dim=tuple(range(1, input.ndim - 1))
          )
      else:
          input_for_gating = input
  
      # 1. compute scores and find most appropriate experts with beam search
      grid_scores = self.proj(input_for_gating).split_with_sizes(
          self.beam_search.grid_size, dim=-1
      )
  
      chosen_experts: List[
          List[RemoteExpert]
      ] = self.beam_search.batch_find_best_experts(
          [scores.detach().cpu().numpy() for scores in grid_scores],
          self.k_best,
      )
  
      if self._expert_info is None:
          try:
              self._expert_info = next(
                  (
                      expert.info
                      for experts_i in chosen_experts
                      for expert in experts_i
                  )
              )
          except StopIteration:
              raise RuntimeError(
                  "No responding experts found during beam search. Check that UID prefixes and "
                  "the grid size are consistent with running Server instances."
              )
          except P2PDaemonError as e:
              logger.warning(
                  f"Failed to get RemoteMixtureOfExperts.output_shape: {e}"
              )
  
      expert_mask, *expert_outputs = _RemoteCallMany.apply(
          DUMMY,
          chosen_experts,
          self.k_min,
          self.backward_k_min,
          self.timeout_after_k_min,
          self.forward_timeout,
          self.backward_timeout,
          self.detect_anomalies,
          self.allow_zero_outputs,
          self.info,
          *nested_flatten(((input, *args), kwargs)),
      )
      # ^-- multiple tensors of shape [batch_size, max_experts, ...output_shape]
  
      expert_logits = self.compute_expert_scores(
          grid_scores, chosen_experts
      )
      masked_logits = torch.full(
          (1,),
          float("-inf"),
          device=expert_logits.device,
          dtype=expert_logits.dtype,
      )
      expert_logits = torch.where(
          expert_mask, expert_logits, masked_logits
      )
      expert_weights = torch.softmax(expert_logits, dim=1)
      averaged_outputs_flat = [
          (expert_weights[..., None] * tensor.flatten(start_dim=2))
          .view(tensor.shape)
          .sum(dim=1)
          for tensor in expert_outputs
      ]  # ^-- multiply by softmax weights along first 2 axes
  
      return nested_pack(
          averaged_outputs_flat, self.info["outputs_schema"]
      )
  
imports: from __future__ import annotations
  import time
  from concurrent.futures import Future
  from queue import Empty, Queue
  from typing import Any, Dict, List, Optional, Tuple
  import torch
  import torch.nn
  from torch.autograd.function import once_differentiable
  from hivemind.compression import serialize_torch_tensor
  from hivemind.dht import DHT
  from hivemind.moe.client.beam_search import MoEBeamSearcher
  from hivemind.moe.client.expert import DUMMY, RemoteExpert, expert_backward, expert_forward, get_server_stub
  from hivemind.moe.client.remote_expert_worker import RemoteExpertWorker
  from hivemind.moe.expert_uid import UID_DELIMITER
  from hivemind.p2p.p2p_daemon_bindings.control import P2PDaemonError
  from hivemind.utils import nested_flatten, nested_map, nested_pack
  from hivemind.utils.logging import get_logger
label_desc: Use image node_id 2 for calling the RemoteMixtureOfExperts obj's underlying member method code with example usage: obj.forward(input) and returns: nested_pack