docstr: Predict masks given image and prompt embeddings.
  
  Args:
      image_embeddings: the embeddings from the image encoder
      image_pe: positional encoding with the shape of image_embeddings
      sparse_prompt_embeddings: the embeddings of the points and boxes
      dense_prompt_embeddings: the embeddings of the mask inputs
      multimask_output: Whether to return multiple masks or a single mask.
  
  Returns:
      batched predicted masks
      batched predictions of mask quality
docstr_tok: ["Predict","masks","given","image","and","prompt","embeddings",".","Args",":","image_embeddings",":","the","embeddings","from","the","image","encoder","image_pe",":","positional","encoding","with","the","shape","of","image_embeddings","sparse_prompt_embeddings",":","the","embeddings","of","the","points","and","boxes","dense_prompt_embeddings",":","the","embeddings","of","the","mask","inputs","multimask_output",":","Whether","to","return","multiple","masks","or","a","single","mask",".","Returns",":","batched","predicted","masks","batched","predictions","of","mask","quality"]
code: def forward(
      self,
      image_embeddings: Tensor,
      image_pe: Tensor,
      sparse_prompt_embeddings: Tensor,
      dense_prompt_embeddings: Tensor,
      multimask_output: bool,
  ) -> tuple[Tensor, Tensor]:
      """Predict masks given image and prompt embeddings.
  
      Args:
          image_embeddings: the embeddings from the image encoder
          image_pe: positional encoding with the shape of image_embeddings
          sparse_prompt_embeddings: the embeddings of the points and boxes
          dense_prompt_embeddings: the embeddings of the mask inputs
          multimask_output: Whether to return multiple masks or a single mask.
  
      Returns:
          batched predicted masks
          batched predictions of mask quality
      """
      masks, iou_pred = self.predict_masks(
          image_embeddings=image_embeddings,
          image_pe=image_pe,
          sparse_prompt_embeddings=sparse_prompt_embeddings,
          dense_prompt_embeddings=dense_prompt_embeddings,
      )
  
      # Select the correct mask or masks for outptu
      if multimask_output:
          mask_slice = slice(1, None)
      else:
          mask_slice = slice(0, 1)
      masks = masks[:, mask_slice, :, :]
      iou_pred = iou_pred[:, mask_slice]
  
      # Prepare output
      return masks, iou_pred
  
file_path: kornia/kornia/contrib/models/sam/architecture/mask_decoder.py
imports: from __future__ import annotations
  import torch
  from torch import nn
  from kornia.contrib.models.common import MLP, LayerNorm2d
  from kornia.core import Module, Tensor, concatenate, stack
label_desc: Use image node_id 2 for calling the MaskDecoder obj's underlying member method code with example usage: obj.forward(image_embeddings, image_pe, sparse_prompt_embeddings, dense_prompt_embeddings, multimask_output) and returns: masks,  iou_pred