code: def test_maxvit_grid_partition(self):
      input_shape = (1, 3, 224, 224)
      partition_size = 7
      n_partitions = input_shape[3] // partition_size
  
      x = torch.randn(input_shape)
      pre_swap = SwapAxes(-2, -3)
      post_swap = SwapAxes(-2, -3)
  
      partition = WindowPartition()
      departition = WindowDepartition()
  
      x_hat = partition(x, n_partitions)
      x_hat = pre_swap(x_hat)
      x_hat = post_swap(x_hat)
      x_hat = departition(
          x_hat, n_partitions, partition_size, partition_size
      )
  
      torch.testing.assert_close(x, x_hat)
  
file_path: vision/test/test_architecture_ops.py
imports: import unittest
  import pytest
  import torch
  from torchvision.models.maxvit import SwapAxes, WindowDepartition, WindowPartition
label_desc: Use image node_id 2 for calling the MaxvitTester obj's underlying member method code with example usage: obj.test_maxvit_grid_partition() without return types