class_docstr: Wrapper for nn.MultiheadAttention which first transposes the input tensor
  from (batch_size, seq_len, feature_dim) to (seq_length, batch_size, feature_dim),
  then applies the attention and transposes the attention outputs back to the input
  shape.
class_docstr_tok: ["Wrapper","for","nn.MultiheadAttention","which","first","transposes","the","input","tensor","from","(","batch_size",",","seq_len",",","feature_dim",")","to","(","seq_length",",","batch_size",",","feature_dim",")",",","then","applies","the","attention","and","transposes","the","attention","outputs","back","to","the","input","shape","."]
docstr: Args:
      x (torch.Tensor): tensor of shape (batch_size, seq_len, feature_dim)
      mask (torch.Tensor): bool tensor with shape (batch_size, seq_len).
          Sequence elements that are False are invalid.
  
  Returns:
      Tensor with shape (batch_size, seq_len, feature_dim)
docstr_tok: ["Args",":","x","(","torch.Tensor",")",":","tensor","of","shape","(","batch_size",",","seq_len",",","feature_dim",")","mask","(","torch.Tensor",")",":","bool","tensor","with","shape","(","batch_size",",","seq_len",")",".","Sequence","elements","that","are","False","are","invalid",".","Returns",":","Tensor","with","shape","(","batch_size",",","seq_len",",","feature_dim",")"]
code: def forward(
      self, x: torch.Tensor, mask: Optional[torch.Tensor] = None
  ) -> torch.Tensor:
      """
      Args:
          x (torch.Tensor): tensor of shape (batch_size, seq_len, feature_dim)
          mask (torch.Tensor): bool tensor with shape (batch_size, seq_len).
              Sequence elements that are False are invalid.
  
      Returns:
          Tensor with shape (batch_size, seq_len, feature_dim)
      """
      assert (
          x.dim() == 3
      ), "Requires x shape (batch_size x seq_len x feature_dim)"
  
      if mask is not None:
          # At least the first element of each masked batch row must be valid for
          # key_padding_mask.
          mask[:, 0] = True
          mask = ~mask
  
      # Transpose x to (seq_length x batch_size x feature_dim).
      x = x.transpose(0, 1)
      attn_output, self._attention_weights = self._attention(
          x, x, x, key_padding_mask=mask
      )
  
      # Transpose attention output to (batch_size x seq_length x feature_dim).
      attn_output = attn_output.transpose(0, 1)
      return attn_output
  
imports: from typing import List, Optional, Tuple
  import torch
  from pytorchvideo.layers.utils import set_attributes
  from torch import nn
  from torch.nn.utils.rnn import pack_padded_sequence
label_desc: Use image node_id 3 for calling the TransposeMultiheadAttention obj's underlying member method code with example usage: obj.forward(x, mask) and returns: attn_output