code: def forward(
      self,
      decoder_embed,
      encoder_embed,
      decoder_pos_embed,
      encoder_pos_embed,
      encoder_k_mask=None,
      decoder_k_mask=None,
      decoder_qk_mask=None,
      return_intermediate=False,
      early_exit_from_layer_idx=None,
  ):
      if early_exit_from_layer_idx:
          assert (
              early_exit_from_layer_idx > 0
              and early_exit_from_layer_idx < self.num_layers
          )
  
      x = decoder_embed
      intermediates = []
      for idx, l in enumerate(self.layers):
          if idx == early_exit_from_layer_idx:
              logger.warning(
                  "Early exit taken from TransformerDecoder "
                  f"after {idx}/{self.num_layers} layers"
              )
              break
  
          x = l(
              x,
              encoder_embed,
              encoder_k_mask,
              decoder_k_mask,
              decoder_qk_mask,
              decoder_pos_embed,
              encoder_pos_embed,
          )
  
          if return_intermediate:
              intermediates.append(self.norm(x))
  
      if return_intermediate:
          return torch.stack(intermediates, 0)
      return self.norm(x)
  
file_path: ml-ane-transformers/ane_transformers/reference/decoder.py
imports: import logging
  import copy
  import torch
  import torch.nn
  from .layer_norm import LayerNormANE
label_desc: Use image node_id 2 for calling the TransformerDecoder obj's underlying member method code with example usage: obj.forward(decoder_embed, encoder_embed, decoder_pos_embed, encoder_pos_embed, encoder_k_mask, decoder_k_mask, decoder_qk_mask, return_intermediate, early_exit_from_layer_idx) and returns: self, torch