class_docstr: ## Routing among multiple FFNs
class_docstr_tok: ["#","#","Routing","among","multiple","FFNs"]
docstr: * `capacity_factor` is the capacity of each expert as a factor relative to ideally balanced load
  * `drop_tokens` specifies whether to drop tokens if more tokens are routed to an expert than the capacity
  * `is_scale_prob` specifies whether to multiply the input to the FFN by the routing probability
  * `n_experts` is the number of experts
  * `expert` is the expert layer, a [FFN module](../feed_forward.html)
  * `d_model` is the number of features in a token embedding
  * `d_ff` is the number of features in the hidden layer of the FFN
  * `dropout` is dropout probability in the FFN
docstr_tok: ["*","`","capacity_factor","`","is","the","capacity","of","each","expert","as","a","factor","relative","to","ideally","balanced","load","*","`","drop_tokens","`","specifies","whether","to","drop","tokens","if","more","tokens","are","routed","to","an","expert","than","the","capacity","*","`","is_scale_prob","`","specifies","whether","to","multiply","the","input","to","the","FFN","by","the","routing","probability","*","`","n_experts","`","is","the","number","of","experts","*","`","expert","`","is","the","expert","layer",",","a","[","FFN","module","]","(","..","\/feed_forward.html",")","*","`","d_model","`","is","the","number","of","features","in","a","token","embedding","*","`","d_ff","`","is","the","number","of","features","in","the","hidden","layer","of","the","FFN","*","`","dropout","`","is","dropout","probability","in","the","FFN"]
code: def __init__(
      self,
      *,
      capacity_factor: float,
      drop_tokens: bool,
      is_scale_prob: bool,
      n_experts: int,
      expert: FeedForward,
      d_model: int,
  ):
      """
      * `capacity_factor` is the capacity of each expert as a factor relative to ideally balanced load
      * `drop_tokens` specifies whether to drop tokens if more tokens are routed to an expert than the capacity
      * `is_scale_prob` specifies whether to multiply the input to the FFN by the routing probability
      * `n_experts` is the number of experts
      * `expert` is the expert layer, a [FFN module](../feed_forward.html)
      * `d_model` is the number of features in a token embedding
      * `d_ff` is the number of features in the hidden layer of the FFN
      * `dropout` is dropout probability in the FFN
      """
      super().__init__()
  
      self.capacity_factor = capacity_factor
      self.is_scale_prob = is_scale_prob
      self.n_experts = n_experts
      self.drop_tokens = drop_tokens
  
      # make copies of the FFNs
      self.experts = clone_module_list(expert, n_experts)
      # Routing layer and softmax
      self.switch = nn.Linear(d_model, n_experts)
      self.softmax = nn.Softmax(dim=-1)
  
file_path: annotated_deep_learning_paper_implementations/labml_nn/transformers/switch/__init__.py
imports: import torch
  from torch import nn
  from labml_helpers.module import Module
  from labml_nn.transformers.feed_forward import FeedForward
  from labml_nn.transformers.mha import MultiHeadAttention
  from labml_nn.utils import clone_module_list
label_desc: Use image node_id 1 to create a new SwitchFeedForward object from inherited base classes: Module with example: obj = SwitchFeedForward()