class_docstr: Layer for partitioning the input tensor into non-overlapping windows and applying attention to each window.
  
  Args:
      in_channels (int): Number of input channels.
      head_dim (int): Dimension of each attention head.
      partition_size (int): Size of the partitions.
      partition_type (str): Type of partitioning to use. Can be either "grid" or "window".
      grid_size (Tuple[int, int]): Size of the grid to partition the input tensor into.
      mlp_ratio (int): Ratio of the  feature size expansion in the MLP layer.
      activation_layer (Callable[..., nn.Module]): Activation function to use.
      norm_layer (Callable[..., nn.Module]): Normalization function to use.
      attention_dropout (float): Dropout probability for the attention layer.
      mlp_dropout (float): Dropout probability for the MLP layer.
      p_stochastic_dropout (float): Probability of dropping out a partition.
docstr: Args:
      x (Tensor): Input tensor with expected layout of [B, C, H, W].
  Returns:
      Tensor: Output tensor with expected layout of [B, C, H, W].
docstr_tok: ["Args",":","x","(","Tensor",")",":","Input","tensor","with","expected","layout","of","[","B",",","C",",","H",",","W","]",".","Returns",":","Tensor",":","Output","tensor","with","expected","layout","of","[","B",",","C",",","H",",","W","]","."]
code: def forward(self, x: Tensor) -> Tensor:
      """
      Args:
          x (Tensor): Input tensor with expected layout of [B, C, H, W].
      Returns:
          Tensor: Output tensor with expected layout of [B, C, H, W].
      """
  
      # Undefined behavior if H or W are not divisible by p
      # https://github.com/google-research/maxvit/blob/da76cf0d8a6ec668cc31b399c4126186da7da944/maxvit/models/maxvit.py#L766
      gh, gw = self.grid_size[0] // self.p, self.grid_size[1] // self.p
      torch._assert(
          self.grid_size[0] % self.p == 0
          and self.grid_size[1] % self.p == 0,
          "Grid size must be divisible by partition size. Got grid size of {} and partition size of {}".format(
              self.grid_size, self.p
          ),
      )
  
      x = self.partition_op(x, self.p)
      x = self.partition_swap(x)
      x = x + self.stochastic_dropout(self.attn_layer(x))
      x = x + self.stochastic_dropout(self.mlp_layer(x))
      x = self.departition_swap(x)
      x = self.departition_op(x, self.p, gh, gw)
  
      return x
  
imports: import math
  from collections import OrderedDict
  from functools import partial
  from typing import Any, Callable, List, Optional, Sequence, Tuple
  import numpy
  import torch
  import torch.nn.functional
  from torch import nn, Tensor
  from torchvision.models._api import register_model, Weights, WeightsEnum
  from torchvision.models._meta import _IMAGENET_CATEGORIES
  from torchvision.models._utils import _ovewrite_named_param, handle_legacy_interface
  from torchvision.ops.misc import Conv2dNormActivation, SqueezeExcitation
  from torchvision.ops.stochastic_depth import StochasticDepth
  from torchvision.transforms._presets import ImageClassification, InterpolationMode
  from torchvision.utils import _log_api_usage_once
label_desc: Use image node_id 2 for calling the PartitionAttentionLayer obj's underlying member method code with example usage: obj.forward(x) and returns: x