class_docstr: Stack of transformer blocks
class_docstr_tok: ["Stack","of","transformer","blocks"]
docstr: In the case of a transformer Decoder: the forward pass predicts the distribution
  of the `latent` given `encoder_output`
  
  Args:
      input: tensor of shape [B', seq_length, hidden_dim]
      encoder_output: tensor of shape [B', seq_length_encoder, hidden_dim],
          which is the result of the encoder output (concatenated).
          Optional (not needed for encoder Transformer), defaults to None.
  
  Returns:
      tensor of shape [B', seq_length, hidden_dim]
docstr_tok: ["In","the","case","of","a","transformer","Decoder",":","the","forward","pass","predicts","the","distribution","of","the","`","latent","`","given","`","encoder_output","`","Args",":","input",":","tensor","of","shape","[","B","'",",","seq_length",",","hidden_dim","]","encoder_output",":","tensor","of","shape","[","B","'",",","seq_length_encoder",",","hidden_dim","]",",","which","is","the","result","of","the","encoder","output","(","concatenated",")",".","Optional","(","not","needed","for","encoder","Transformer",")",",","defaults","to","None",".","Returns",":","tensor","of","shape","[","B","'",",","seq_length",",","hidden_dim","]"]
code: def forward(
      self, input: Tensor, encoder_output: Optional[Tensor]
  ) -> Tensor:
      """
      In the case of a transformer Decoder: the forward pass predicts the distribution
      of the `latent` given `encoder_output`
  
      Args:
          input: tensor of shape [B', seq_length, hidden_dim]
          encoder_output: tensor of shape [B', seq_length_encoder, hidden_dim],
              which is the result of the encoder output (concatenated).
              Optional (not needed for encoder Transformer), defaults to None.
  
      Returns:
          tensor of shape [B', seq_length, hidden_dim]
      """
      assert (
          input.dim() == 3
      ), f"Expected (batch_size, seq_length, hidden_dim) got {input.shape}"
      if encoder_output is not None:
          if input.shape[-1] != encoder_output.shape[-1]:
              raise ValueError(
                  f"Expected latent.shape[-1] == encoder_output.shape[-1], "
                  f"got {input.shape[-1]}, {encoder_output.shape[-1]}"
              )
  
      for layer in self.layers:
          input = layer(input=input, encoder_output=encoder_output)
      return input
  
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 Transformer obj's underlying member method code with example usage: obj.forward(input, encoder_output) and returns: input