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.
code: def __init__(
      self,
      in_channels: int,
      head_dim: int,
      # partitioning parameters
      partition_size: int,
      partition_type: str,
      # grid size needs to be known at initialization time
      # because we need to know hamy relative offsets there are in the grid
      grid_size: Tuple[int, int],
      mlp_ratio: int,
      activation_layer: Callable[..., nn.Module],
      norm_layer: Callable[..., nn.Module],
      attention_dropout: float,
      mlp_dropout: float,
      p_stochastic_dropout: float,
  ) -> None:
      super().__init__()
  
      self.n_heads = in_channels // head_dim
      self.head_dim = head_dim
      self.n_partitions = grid_size[0] // partition_size
      self.partition_type = partition_type
      self.grid_size = grid_size
  
      if partition_type not in ["grid", "window"]:
          raise ValueError(
              "partition_type must be either 'grid' or 'window'"
          )
  
      if partition_type == "window":
          self.p, self.g = partition_size, self.n_partitions
      else:
          self.p, self.g = self.n_partitions, partition_size
  
      self.partition_op = WindowPartition()
      self.departition_op = WindowDepartition()
      self.partition_swap = (
          SwapAxes(-2, -3)
          if partition_type == "grid"
          else nn.Identity()
      )
      self.departition_swap = (
          SwapAxes(-2, -3)
          if partition_type == "grid"
          else nn.Identity()
      )
  
      self.attn_layer = nn.Sequential(
          norm_layer(in_channels),
          # it's always going to be partition_size ** 2 because
          # of the axis swap in the case of grid partitioning
          RelativePositionalMultiHeadAttention(
              in_channels, head_dim, partition_size**2
          ),
          nn.Dropout(attention_dropout),
      )
  
      # pre-normalization similar to transformer layers
      self.mlp_layer = nn.Sequential(
          nn.LayerNorm(in_channels),
          nn.Linear(in_channels, in_channels * mlp_ratio),
          activation_layer(),
          nn.Linear(in_channels * mlp_ratio, in_channels),
          nn.Dropout(mlp_dropout),
      )
  
      # layer scale factors
      self.stochastic_dropout = StochasticDepth(
          p_stochastic_dropout, mode="row"
      )
  
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 1 to create a new PartitionAttentionLayer object from inherited base classes: nn with example: obj = PartitionAttentionLayer(in_channels, head_dim, partition_size, partition_type, grid_size, mlp_ratio, activation_layer, norm_layer, attention_dropout, mlp_dropout, p_stochastic_dropout)