class_docstr: A module implementing Switch Transformers [1] Mixture-of-Experts inference with remote experts.
  
  [1] Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity.
   William Fedus, Barret Zoph, Noam Shazeer. https://arxiv.org/abs/2101.03961
  
  :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 just the input if no experts respond on forward pass
class_docstr_tok: ["A","module","implementing","Switch","Transformers","[","1","]","Mixture-of-Experts","inference","with","remote","experts",".","[","1","]","Switch","Transformers",":","Scaling","to","Trillion","Parameter","Models","with","Simple","and","Efficient","Sparsity",".","William","Fedus",",","Barret","Zoph",",","Noam","Shazeer",".","https",":","\/\/arxiv.org\/abs\/2101.03961",":","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","just","the","input","if","no","experts","respond","on","forward","pass"]
code: def __init__(
      self,
      *,
      grid_size: Tuple[int, ...],
      utilization_alpha: float = 0.9,
      grid_dropout: float = 1.0,
      jitter_eps: float = 1e-2,
      k_best=1,
      k_min=0,
      backward_k_min=0,
      allow_zero_outputs=True,
      **kwargs,
  ):
      super().__init__(
          grid_size=grid_size,
          k_best=k_best,
          k_min=k_min,
          backward_k_min=backward_k_min,
          allow_zero_outputs=allow_zero_outputs,
          **kwargs,
      )
  
      initial_utilization = torch.cat(
          [
              torch.tensor(
                  [1 / dim_size for _ in range(dim_size)],
                  dtype=torch.float,
              )
              for dim_size in grid_size
          ],
      )
      self.register_buffer("grid_utilization", initial_utilization)
      self.utilization_alpha = utilization_alpha
      self.grid_dropout = grid_dropout
      self.jitter_eps = jitter_eps
  
file_path: hivemind/hivemind/moe/client/switch_moe.py
imports: from __future__ import annotations
  from typing import List, Tuple
  import torch
  from hivemind.moe.client.expert import DUMMY, RemoteExpert
  from hivemind.moe.client.moe import RemoteMixtureOfExperts, _RemoteCallMany
  from hivemind.moe.expert_uid import UID_DELIMITER
  from hivemind.p2p.p2p_daemon_bindings.control import P2PDaemonError
  from hivemind.utils import nested_flatten, nested_pack
  from hivemind.utils.logging import get_logger
label_desc: Use image node_id 1 to create a new RemoteSwitchMixtureOfExperts object from inherited base classes: RemoteMixtureOfExperts with example: obj = RemoteSwitchMixtureOfExperts()