code: def window_partition(
      t: torch.Tensor, window_size: int
  ) -> torch.Tensor:
      """
      Partition tensor @t into chunks of size @window_size.
  
      @t's sequence length must be divisible by @window_size.
  
      Args:
          t: A tensor of shape [batch_size, sequence_length, embed_dim].
          window_size: The desired window size.
  
      Returns:
          A tensor of shape [batch_size * sequence_length // window_size,
          window_size, embed_dim].
      """
      B, N, C = t.shape
  
      if not N % window_size == 0:
          raise ValueError(
              f"sequence length {N} must be divisible by window size {window_size}"
          )
  
      t = t.reshape(B * N // window_size, window_size, C)
      return t
  
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 1 for calling a global function with example usage: window_partition(t, window_size) and returns: t