code: def pytorch_decoder_mask(
      memory: torch.Tensor, tgt_in_idx: torch.Tensor, num_heads: int
  ):
      """
      Compute the masks used in the PyTorch Transformer-based decoder for
      self-attention and attention over encoder outputs
  
      mask_ijk = 1 if the item should be ignored; 0 if the item should be paid attention
  
      Input:
          memory shape: batch_size, src_seq_len, dim_model
          tgt_in_idx (+2 offseted) shape: batch_size, tgt_seq_len
  
      Return:
          tgt_tgt_mask shape: batch_size * num_heads, tgt_seq_len, tgt_seq_len
          tgt_src_mask shape: batch_size * num_heads, tgt_seq_len, src_seq_len
      """
      batch_size, src_seq_len, _ = memory.shape
      tgt_seq_len = tgt_in_idx.shape[1]
      device = memory.device
      mask_indices = torch.tril(
          tgt_in_idx.repeat(1, tgt_seq_len).reshape(
              batch_size, tgt_seq_len, tgt_seq_len
          ),
          diagonal=0,
      ).to(device)
      tgt_src_mask_augmented = torch.zeros(
          batch_size,
          tgt_seq_len,
          src_seq_len + 2,
          dtype=torch.bool,
          device=device,
      ).scatter(2, mask_indices, 1)
      tgt_src_mask = tgt_src_mask_augmented[:, :, 2:].repeat_interleave(
          num_heads, dim=0
      )
      tgt_tgt_mask = (subsequent_mask(tgt_seq_len, device) == 0).repeat(
          batch_size * num_heads, 1, 1
      )
      return tgt_tgt_mask, tgt_src_mask
  
imports: import copy
  import logging
  import math
  from enum import Enum
  import torch
  import torch.nn
  import torch.nn.functional
label_desc: Use image node_id 9 for calling a global function with example usage: pytorch_decoder_mask(memory, tgt_in_idx, num_heads) and returns: tgt_tgt_mask,  tgt_src_mask