class_docstr: This class defines the pre-norm `Transformer encoder <https://arxiv.org/abs/1706.03762>`_
  with the addition of windowed attention.
  
  This class first partitions the input sequence into a series of windows (with
  an optional offset to use when defining windows). Then, it calls a
  TransformerEncoder module. Then, it undoes windowing.
  
  Args:
      opts: Command line arguments.
      embed_dim: :math:`C_{in}` from an expected input of size :math:`(N, P, C_{in})`.
      ffn_latent_dim: Inner dimension of the FFN.
      num_heads: Number of heads in multi-head attention. Default: 8.
      attn_dropout: Dropout rate for attention in multi-head attention. Default: 0.0.
      dropout: Dropout rate. Default: 0.0.
      ffn_dropout: Dropout between FFN layers. Default: 0.0.
      transformer_norm_layer: Normalization layer. Default: layer_norm.
      stochastic_dropout: Stochastic dropout setting. Default: 0.0.
      window_size: The size of the window, if using windowed attention. Default: None.
      window_shift: The size of the shift, if using shifted windowed attention. Default: None.
  
  Shape:
      - Input: :math:`(N, P, C_{in})` where :math:`N` is batch size, :math:`P` is number of patches,
      and :math:`C_{in}` is input embedding dim
      - Output: same shape as the input
docstr: Compute the outputs of the WindowedTransformerEncoder on an input.
  
  Args:
      x: The input tensor, of shape [batch_size, sequence_length, embed_dim].
      x_prev: The context input, if using cross-attention. Its shape is
          [batch_size, sequence_length_2, embed_dim].
      key_padding_mask: An optional tensor of masks to be applied to the
          inputs @x. Its shape is [batch_size, sequence_length].
      attn_mask: An optional attention mask. Its shape is [batch_size,
          sequence_length, sequence_length_2]. (If using self-attention,
          the sequence lengths will be equal.)
  
  Returns:
      The WindowedTransformerEncoder output.
docstr_tok: ["Compute","the","outputs","of","the","WindowedTransformerEncoder","on","an","input",".","Args",":","x",":","The","input","tensor",",","of","shape","[","batch_size",",","sequence_length",",","embed_dim","]",".","x_prev",":","The","context","input",",","if","using","cross-attention",".","Its","shape","is","[","batch_size",",","sequence_length_2",",","embed_dim","]",".","key_padding_mask",":","An","optional","tensor","of","masks","to","be","applied","to","the","inputs","@","x",".","Its","shape","is","[","batch_size",",","sequence_length","]",".","attn_mask",":","An","optional","attention","mask",".","Its","shape","is","[","batch_size",",","sequence_length",",","sequence_length_2","]",".","(","If","using","self-attention",",","the","sequence","lengths","will","be","equal",".",")","Returns",":","The","WindowedTransformerEncoder","output","."]
code: def forward(
      self,
      x: Tensor,
      x_prev: Optional[Tensor] = None,
      key_padding_mask: Optional[Tensor] = None,
      attn_mask: Optional[Tensor] = None,
      *args,
      **kwargs,
  ) -> Tensor:
      """
      Compute the outputs of the WindowedTransformerEncoder on an input.
  
      Args:
          x: The input tensor, of shape [batch_size, sequence_length, embed_dim].
          x_prev: The context input, if using cross-attention. Its shape is
              [batch_size, sequence_length_2, embed_dim].
          key_padding_mask: An optional tensor of masks to be applied to the
              inputs @x. Its shape is [batch_size, sequence_length].
          attn_mask: An optional attention mask. Its shape is [batch_size,
              sequence_length, sequence_length_2]. (If using self-attention,
              the sequence lengths will be equal.)
  
      Returns:
          The WindowedTransformerEncoder output.
      """
      B, N, C = x.shape
      (
          x,
          windowed_key_padding_mask,
          windows_mask,
      ) = window_x_and_key_padding_mask(
          x, key_padding_mask, self.window_size, self.window_shift
      )
      total_mask = windowed_key_padding_mask.unsqueeze(1) + windows_mask
  
      if attn_mask is not None:
          total_mask += attn_mask
  
      # If an entire window is masked out, attention is computed across
      # only -inf values, which gives NaN. We instead set these masks to
      # 0 to avoid this.
      fully_masked_windows = total_mask.max(dim=-1).values == float(
          "-inf"
      )
      total_mask[fully_masked_windows] = 0
  
      x = super().forward(x, x_prev, attn_mask=attn_mask)
  
      # Undo windowing.
      x = unwindow_x(x, B, N, C, self.window_shift)
      return x
  
file_path: ml-cvnets/cvnets/modules/windowed_transformer.py
imports: import argparse
  from typing import Optional, Tuple
  import torch
  from torch import Tensor
  from cvnets.layers import token_merging
  from cvnets.modules import transformer
label_desc: Use image node_id 2 for calling the WindowedTransformerEncoder obj's underlying member method code with example usage: obj.forward(x, x_prev, key_padding_mask, attn_mask) and returns: x