class_docstr: # Switch Transformer Block
  
  This is the same as [normal transformer block](../models.html#TransformerLayer)
  with handling extra outputs of switch feedforward module.
class_docstr_tok: ["#","Switch","Transformer","Block","This","is","the","same","as","[","normal","transformer","block","]","(","..","\/models.html","#","TransformerLayer",")","with","handling","extra","outputs","of","switch","feedforward","module","."]
code: def forward(self, *, x: torch.Tensor, mask: torch.Tensor):
      # Normalize the vectors before doing self attention
      z = self.norm_self_attn(x)
      # Run through self attention, i.e. keys and values are from self
      self_attn = self.attn(query=z, key=z, value=z, mask=mask)
      # Add the self attention results
      x = x + self.dropout(self_attn)
  
      # Normalize for feed-forward
      z = self.norm_ff(x)
      # Pass through the switching feed-forward network
      (
          ff,
          counts,
          route_prob,
          n_dropped,
          route_prob_max,
      ) = self.feed_forward(z)
      # Add the feed-forward results back
      x = x + self.dropout(ff)
  
      return x, counts, route_prob, n_dropped, route_prob_max
  
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 2 for calling the SwitchTransformerLayer obj's underlying member method code with example usage: obj.forward() and returns: x,  counts,  route_prob,  n_dropped,  route_prob_max