code: def forward(self, x, mask_features, extra=None):
      if self.task == "caption":
          return self.forward_caption(x, mask_features, extra)
  
      assert len(x) == self.num_feature_levels
      src = []
      pos = []
      size_list = []
  
      for i in range(self.num_feature_levels):
          size_list.append(x[i].shape[-2:])
          pos.append(self.pe_layer(x[i], None).flatten(2))
          src.append(
              self.input_proj[i](x[i]).flatten(2)
              + self.level_embed.weight[i][None, :, None]
          )
  
          # flatten NxCxHxW to HWxNxC
          pos[-1] = pos[-1].permute(2, 0, 1)
          src[-1] = src[-1].permute(2, 0, 1)
  
      _, bs, _ = src[0].shape
  
      query_embed = self.query_embed.weight.unsqueeze(1).repeat(
          1, bs, 1
      )
      output = self.query_feat.weight.unsqueeze(1).repeat(1, bs, 1)
  
      predictions_mask = []
      predictions_class_embed = []
  
      if self.task == "ref-seg":
          self_tgt_mask = self.self_attn_mask[
              :, : self.num_queries, : self.num_queries
          ].repeat(output.shape[1] * self.num_heads, 1, 1)
          grounding_tokens = extra["grounding_tokens"]
          _grounding_tokens = grounding_tokens.detach().clone()
          # initialize with negative attention at the beginning.
          pad_tgt_mask = (
              torch.ones(
                  (
                      1,
                      self.num_queries
                      + (self.num_queries - 1)
                      + len(grounding_tokens),
                      self.num_queries
                      + (self.num_queries - 1)
                      + len(grounding_tokens),
                  ),
                  device=self_tgt_mask.device,
              )
              .bool()
              .repeat(output.shape[1] * self.num_heads, 1, 1)
          )
          pad_tgt_mask[
              :, : self.num_queries, : self.num_queries
          ] = self_tgt_mask
          # grounding tokens could attend with eatch other
          pad_tgt_mask[
              :, self.num_queries :, self.num_queries :
          ] = False
          self_tgt_mask = pad_tgt_mask
          output = torch.cat((output, output[:-1]), dim=0)
          # also pad language embdding to fix embedding
          query_embed = torch.cat(
              (query_embed, query_embed[:-1]), dim=0
          )
      else:
          self_tgt_mask = self.self_attn_mask[
              :, : self.num_queries, : self.num_queries
          ].repeat(output.shape[1] * self.num_heads, 1, 1)
  
      results = self.forward_prediction_heads(
          output, mask_features, attn_mask_target_size=size_list[0]
      )
      attn_mask = results["attn_mask"]
      predictions_class_embed.append(results["class_embed"])
      predictions_mask.append(results["outputs_mask"])
  
      for i in range(self.num_layers):
          level_index = i % self.num_feature_levels
          attn_mask[
              torch.where(attn_mask.sum(-1) == attn_mask.shape[-1])
          ] = False
  
          # attention: cross-attention first
          output, avg_attn = self.transformer_cross_attention_layers[i](
              output,
              src[level_index],
              memory_mask=attn_mask,
              # here we do not apply masking on padded region
              memory_key_padding_mask=None,
              pos=pos[level_index],
              query_pos=query_embed,
          )
  
          if self.task == "ref-seg":
              output = torch.cat((output, _grounding_tokens), dim=0)
              query_embed = torch.cat(
                  (query_embed, grounding_tokens), dim=0
              )
  
          output = self.transformer_self_attention_layers[i](
              output,
              tgt_mask=self_tgt_mask,
              tgt_key_padding_mask=None,
              query_pos=query_embed,
          )
  
          output = self.transformer_ffn_layers[i](output)
  
          if self.task == "ref-seg":
              _grounding_tokens = output[-len(_grounding_tokens) :]
              output = output[: -len(_grounding_tokens)]
              query_embed = query_embed[: -len(_grounding_tokens)]
  
          results = self.forward_prediction_heads(
              output,
              mask_features,
              attn_mask_target_size=size_list[
                  (i + 1) % self.num_feature_levels
              ],
          )
          attn_mask = results["attn_mask"]
          predictions_mask.append(results["outputs_mask"])
          predictions_class_embed.append(results["class_embed"])
  
      out = {
          "pred_masks": predictions_mask[-1],
          "pred_class_embed": predictions_class_embed[-1],
      }
  
      if self.task == "ref-seg":
          mask_pred_results = []
          outputs_class = []
          for idx in range(mask_features.shape[0]):  # batch size
              pred_gmasks = out["pred_masks"][
                  idx, self.num_queries : 2 * self.num_queries - 1
              ]
              v_emb = predictions_class_embed[-1][
                  idx, self.num_queries : 2 * self.num_queries - 1
              ]
              t_emb = extra["class_emb"]
  
              t_emb = t_emb / (t_emb.norm(dim=-1, keepdim=True) + 1e-7)
              v_emb = v_emb / (v_emb.norm(dim=-1, keepdim=True) + 1e-7)
  
              temperature = self.lang_encoder.logit_scale
              out_prob = vl_similarity(
                  v_emb, t_emb, temperature=temperature
              )
  
              matched_id = out_prob.max(0)[1]
              mask_pred_results += [pred_gmasks[matched_id, :, :]]
              outputs_class += [out_prob[matched_id, :]]
          out["pred_masks"] = mask_pred_results
          out["pred_logits"] = outputs_class
      elif self.task == "retrieval":
          t_emb = extra["class_emb"]
          temperature = self.lang_encoder.logit_scale
          v_emb = out["pred_class_embed"][:, -1, :]
          v_emb = v_emb / (v_emb.norm(dim=-1, keepdim=True) + 1e-7)
          logits = vl_similarity(v_emb, t_emb, temperature)
          out["pred_logits"] = logits
      elif self.task in ["semseg", "instance", "panoptic"]:
          outputs_class = self.lang_encoder.compute_similarity(
              out["pred_class_embed"]
          )
          out["pred_logits"] = outputs_class
      return out
  
file_path: mmdetection/projects/XDecoder/xdecoder/transformer_decoder.py
imports: import torch
  from torch import nn
  from torch.nn import functional
  from mmdet.registry import MODELS
  from .language_model import LanguageEncoder
  from .transformer_blocks import MLP, Conv2d, CrossAttentionLayer, FFNLayer, PositionEmbeddingSine, SelfAttentionLayer
  from .utils import is_lower_torch_version
label_desc: Use image node_id 2 for calling the XDecoderTransformerDecoder obj's underlying member method code with example usage: obj.forward(x, mask_features, extra) and returns: out, self