class_docstr: Merge tokens from a [batch_size, sequence_length, num_channels] tensor
  using a linear projection.
  
  This function also updates masks and adds padding as needed to make the
  sequence length divisible by the window size before merging tokens.
  
  Args:
      dim: Number of input channels.
      window: The size of the window to merge into a single token.
docstr: Perform token merging.
  
  Args:
      x: A tensor of shape [batch_size, sequence_length, num_channels].
      key_padding_mask: A tensor of shape [batch_size, sequence_length]
          with "-inf" values at mask tokens, and "0" values at unmasked
          tokens.
  
  Returns:
      A tensor of shape [batch_size, math.ceil(sequence_length /
          self.window), num_channels], where @self.window is the window
          size.
code: def forward(
      self, x: torch.Tensor, key_padding_mask: torch.Tensor
  ) -> Tuple[torch.Tensor, torch.Tensor]:
      """
      Perform token merging.
  
      Args:
          x: A tensor of shape [batch_size, sequence_length, num_channels].
          key_padding_mask: A tensor of shape [batch_size, sequence_length]
              with "-inf" values at mask tokens, and "0" values at unmasked
              tokens.
  
      Returns:
          A tensor of shape [batch_size, math.ceil(sequence_length /
              self.window), num_channels], where @self.window is the window
              size.
      """
      if key_padding_mask is not None:
          # Zero out the masked portion of @x to make sure it doesn't
          # participate in linear projections after windowing.
          x[key_padding_mask == float("-inf")] = 0
  
      x, key_padding_mask = pad_x_and_mask(
          x, key_padding_mask, self.window
      )
      B, N, C = x.shape
  
      x = x.unfold(
          1, self.window, self.window
      )  # [B, N // window, C, window]
      x = x.reshape(B, N // self.window, C * self.window)
      x = self.reduction(x)  # [B, N // self.window, C]
      x = self.norm(x)
  
      if key_padding_mask is not None:
          assert key_padding_mask.shape == (B, N)
          key_padding_mask = key_padding_mask.unfold(
              1, self.window, self.window
          )  # [B, N // window, window].
          key_padding_mask = key_padding_mask.max(
              dim=-1
          ).values  # [B, N // window].
  
      return x, key_padding_mask
  
file_path: ml-cvnets/cvnets/layers/token_merging.py
imports: from typing import Tuple
  import torch
  from torch import nn
  from torch.nn import functional
  from cvnets.layers import linear_layer
  from cvnets.layers.normalization import layer_norm
label_desc: Use image node_id 2 for calling the TokenMerging obj's underlying member method code with example usage: obj.forward(x, key_padding_mask) and returns: x,  key_padding_mask