class_docstr: Single transformer block, can be used for both encoder and decoder
class_docstr_tok: ["Single","transformer","block",",","can","be","used","for","both","encoder","and","decoder"]
docstr: _summary_
  
  Args:
      input: tensor with expected shape [B, seq_len, C]
      encoder_output: output from the encoder, used in decoder only.
          Defaults to None.
  
  Returns:
      tensor of the same shape of `input`
docstr_tok: ["_summary_","Args",":","input",":","tensor","with","expected","shape","[","B",",","seq_len",",","C","]","encoder_output",":","output","from","the","encoder",",","used","in","decoder","only",".","Defaults","to","None",".","Returns",":","tensor","of","the","same","shape","of","`","input","`"]
code: def forward(
      self,
      input: Tensor,
      encoder_output: Optional[Tensor] = None,
  ) -> Tensor:
      """_summary_
  
      Args:
          input: tensor with expected shape [B, seq_len, C]
          encoder_output: output from the encoder, used in decoder only.
              Defaults to None.
  
      Returns:
          tensor of the same shape of `input`
      """
      assert (
          input.dim() == 3
      ), f"Expected (batch_size, seq_length, hidden_dim) got {input.shape}"
      if self.is_decoder and encoder_output is None:
          raise ValueError("Decoder needs `encoder_output`")
  
      # ------ FIRST BLOCK ------
      x = self.ln_1a(input)  # normalise input
      x, _ = self.self_attention(
          query=x, key=x, value=x, mask=self.look_ahead_mask
      )  # attend
      x = input + self.drop_path(x)  # skip connnection
      # normalise -> MLP -> add to itself == OUTPUT of FIRST BLOCK
      x = x + self.drop_path(
          self.mlp1(self.ln_1b(x))
      )  # [B, seq_len, hidden_dim]
  
      # ------ SECOND BLOCK ------
      y = self.ln_2a(x)  # normalise input, [B, seq_len, hidden_dim]
      y, _ = self.cross_attention(
          query=y,
          key=encoder_output if encoder_output is not None else y,
          value=encoder_output if encoder_output is not None else y,
          mask=None,
      )  # attend
      y = x + self.drop_path(y)  # skip connection
      # normalise -> MLP -> add to itself == OUTPUT of SECOND BLOCK
      y = y + self.drop_path(
          self.mlp2(self.ln_2b(y))
      )  # [B, seq_len, hidden_dim]
      return y
  
imports: from functools import partial
  from typing import Callable, Optional
  import torch
  import torch.nn
  from torch import Tensor
  from .layers_utils import MLP, StochasticDepth, WindowMultiHeadAttention
label_desc: Use image node_id 2 for calling the TransformerBlock obj's underlying member method code with example usage: obj.forward(input, encoder_output) and returns: y