docstr: Pass the input through the encoder layer.
  
  Args:
      hidden_states: the sequence to the encoder layer (required).
      residual: if postnorm, residual=None, If prenorm, hidden_states = LayerNorm(residual)
docstr_tok: ["Pass","the","input","through","the","encoder","layer",".","Args",":","hidden_states",":","the","sequence","to","the","encoder","layer","(","required",")",".","residual",":","if","postnorm",",","residual=None",",","If","prenorm",",","hidden_states","=","LayerNorm","(","residual",")"]
code: def forward(
      self, hidden_states: Tensor, residual: Optional[Tensor] = None
  ):
      r"""Pass the input through the encoder layer.
  
      Args:
          hidden_states: the sequence to the encoder layer (required).
          residual: if postnorm, residual=None, If prenorm, hidden_states = LayerNorm(residual)
      """
      if self.prenorm:
          assert residual is not None
          mixer_out = self.mixer(hidden_states)
          if not self.fused_dropout_add_ln:
              residual = (
                  self.drop_path1(self.dropout1(mixer_out)) + residual
              )
              hidden_states = self.norm1(
                  residual.to(dtype=self.norm1.weight.dtype)
              )
          else:
              if self.drop_path1.p == 0 or not self.training:
                  rowscale1 = None
              else:
                  rowscale1 = self.drop_path1(
                      torch.ones(
                          mixer_out.shape[:-1],
                          device=mixer_out.device,
                          dtype=mixer_out.dtype,
                      )
                  )
              hidden_states, residual = dropout_add_layer_norm(
                  mixer_out,
                  residual,
                  self.norm1.weight,
                  self.norm1.bias,
                  self.dropout1.p if self.training else 0.0,
                  self.norm1.eps,
                  rowscale=rowscale1,
                  prenorm=True,
              )
          if not isinstance(self.mlp, nn.Identity):
              mlp_out = self.mlp(hidden_states)
              if not self.fused_dropout_add_ln:
                  residual = (
                      self.drop_path2(self.dropout2(mlp_out)) + residual
                  )
                  hidden_states = self.norm2(
                      residual.to(dtype=self.norm2.weight.dtype)
                  )
              else:
                  if self.drop_path2.p == 0 or not self.training:
                      rowscale2 = None
                  else:
                      rowscale2 = self.drop_path2(
                          torch.ones(
                              mlp_out.shape[:-1],
                              device=mlp_out.device,
                              dtype=mlp_out.dtype,
                          )
                      )
                  hidden_states, residual = dropout_add_layer_norm(
                      mlp_out,
                      residual,
                      self.norm2.weight,
                      self.norm2.bias,
                      self.dropout2.p if self.training else 0.0,
                      self.norm2.eps,
                      rowscale=rowscale2,
                      prenorm=True,
                  )
          return hidden_states, residual
      else:
          assert residual is None
          mixer_out = self.mixer(hidden_states)
          if not self.fused_dropout_add_ln:
              hidden_states = self.norm1(
                  (
                      self.drop_path1(self.dropout1(mixer_out))
                      + hidden_states
                  ).to(dtype=self.norm1.weight.dtype)
              )
          else:
              if self.drop_path1.p == 0 or not self.training:
                  rowscale1 = None
              else:
                  rowscale1 = self.drop_path1(
                      torch.ones(
                          mixer_out.shape[:-1],
                          device=mixer_out.device,
                          dtype=mixer_out.dtype,
                      )
                  )
              hidden_states = dropout_add_layer_norm(
                  mixer_out,
                  hidden_states,
                  self.norm1.weight,
                  self.norm1.bias,
                  self.dropout1.p if self.training else 0.0,
                  self.norm1.eps,
                  rowscale=rowscale1,
                  prenorm=False,
              )
          if not isinstance(self.mlp, nn.Identity):
              mlp_out = self.mlp(hidden_states)
              if not self.fused_dropout_add_ln:
                  hidden_states = self.norm2(
                      (
                          self.drop_path2(self.dropout2(mlp_out))
                          + hidden_states
                      ).to(dtype=self.norm2.weight.dtype)
                  )
              else:
                  if self.drop_path2.p == 0 or not self.training:
                      rowscale2 = None
                  else:
                      rowscale2 = self.drop_path2(
                          torch.ones(
                              mlp_out.shape[:-1],
                              device=mlp_out.device,
                              dtype=mlp_out.dtype,
                          )
                      )
                  hidden_states = dropout_add_layer_norm(
                      mlp_out,
                      hidden_states,
                      self.norm2.weight,
                      self.norm2.bias,
                      self.dropout2.p if self.training else 0.0,
                      self.norm2.eps,
                      rowscale=rowscale2,
                      prenorm=False,
                  )
          return hidden_states
  
file_path: zoo/src/models/modules/block.py
imports: from typing import Optional
  from functools import partial
  import torch
  import torch.nn
  import torch.nn.functional
  from torch import Tensor
  from torchvision.ops import StochasticDepth
  from src.models.modules.mha import MHA
  from src.models.modules.mlp import Mlp
label_desc: Use image node_id 2 for calling the Block obj's underlying member method code with example usage: obj.forward(hidden_states, residual) and returns: hidden_states,  residual, hidden_states