class_docstr: Implements MaxVit Transformer from the `MaxViT: Multi-Axis Vision Transformer <https://arxiv.org/abs/2204.01697>`_ paper.
  Args:
      input_size (Tuple[int, int]): Size of the input image.
      stem_channels (int): Number of channels in the stem.
      partition_size (int): Size of the partitions.
      block_channels (List[int]): Number of channels in each block.
      block_layers (List[int]): Number of layers in each block.
      stochastic_depth_prob (float): Probability of stochastic depth. Expands to a list of probabilities for each layer that scales linearly to the specified value.
      squeeze_ratio (float): Squeeze ratio in the SE Layer. Default: 0.25.
      expansion_ratio (float): Expansion ratio in the MBConv bottleneck. Default: 4.
      norm_layer (Callable[..., nn.Module]): Normalization function. Default: None (setting to None will produce a `BatchNorm2d(eps=1e-3, momentum=0.99)`).
      activation_layer (Callable[..., nn.Module]): Activation function Default: nn.GELU.
      head_dim (int): Dimension of the attention heads.
      mlp_ratio (int): Expansion ratio of the MLP layer. Default: 4.
      mlp_dropout (float): Dropout probability for the MLP layer. Default: 0.0.
      attention_dropout (float): Dropout probability for the attention layer. Default: 0.0.
      num_classes (int): Number of classes. Default: 1000.
code: def __init__(
      self,
      # input size parameters
      input_size: Tuple[int, int],
      # stem and task parameters
      stem_channels: int,
      # partitioning parameters
      partition_size: int,
      # block parameters
      block_channels: List[int],
      block_layers: List[int],
      # attention head dimensions
      head_dim: int,
      stochastic_depth_prob: float,
      # conv + transformer parameters
      # norm_layer is applied only to the conv layers
      # activation_layer is applied both to conv and transformer layers
      norm_layer: Optional[Callable[..., nn.Module]] = None,
      activation_layer: Callable[..., nn.Module] = nn.GELU,
      # conv parameters
      squeeze_ratio: float = 0.25,
      expansion_ratio: float = 4,
      # transformer parameters
      mlp_ratio: int = 4,
      mlp_dropout: float = 0.0,
      attention_dropout: float = 0.0,
      # task parameters
      num_classes: int = 1000,
  ) -> None:
      super().__init__()
      _log_api_usage_once(self)
  
      input_channels = 3
  
      # https://github.com/google-research/maxvit/blob/da76cf0d8a6ec668cc31b399c4126186da7da944/maxvit/models/maxvit.py#L1029-L1030
      # for the exact parameters used in batchnorm
      if norm_layer is None:
          norm_layer = partial(nn.BatchNorm2d, eps=1e-3, momentum=0.99)
  
      # Make sure input size will be divisible by the partition size in all blocks
      # Undefined behavior if H or W are not divisible by p
      # https://github.com/google-research/maxvit/blob/da76cf0d8a6ec668cc31b399c4126186da7da944/maxvit/models/maxvit.py#L766
      block_input_sizes = _make_block_input_shapes(
          input_size, len(block_channels)
      )
      for idx, block_input_size in enumerate(block_input_sizes):
          if (
              block_input_size[0] % partition_size != 0
              or block_input_size[1] % partition_size != 0
          ):
              raise ValueError(
                  f"Input size {block_input_size} of block {idx} is not divisible by partition size {partition_size}. "
                  f"Consider changing the partition size or the input size.\n"
                  f"Current configuration yields the following block input sizes: {block_input_sizes}."
              )
  
      # stem
      self.stem = nn.Sequential(
          Conv2dNormActivation(
              input_channels,
              stem_channels,
              3,
              stride=2,
              norm_layer=norm_layer,
              activation_layer=activation_layer,
              bias=False,
              inplace=None,
          ),
          Conv2dNormActivation(
              stem_channels,
              stem_channels,
              3,
              stride=1,
              norm_layer=None,
              activation_layer=None,
              bias=True,
          ),
      )
  
      # account for stem stride
      input_size = _get_conv_output_shape(
          input_size, kernel_size=3, stride=2, padding=1
      )
      self.partition_size = partition_size
  
      # blocks
      self.blocks = nn.ModuleList()
      in_channels = [stem_channels] + block_channels[:-1]
      out_channels = block_channels
  
      # precompute the stochastich depth probabilities from 0 to stochastic_depth_prob
      # since we have N blocks with L layers, we will have N * L probabilities uniformly distributed
      # over the range [0, stochastic_depth_prob]
      p_stochastic = np.linspace(
          0, stochastic_depth_prob, sum(block_layers)
      ).tolist()
  
      p_idx = 0
      for in_channel, out_channel, num_layers in zip(
          in_channels, out_channels, block_layers
      ):
          self.blocks.append(
              MaxVitBlock(
                  in_channels=in_channel,
                  out_channels=out_channel,
                  squeeze_ratio=squeeze_ratio,
                  expansion_ratio=expansion_ratio,
                  norm_layer=norm_layer,
                  activation_layer=activation_layer,
                  head_dim=head_dim,
                  mlp_ratio=mlp_ratio,
                  mlp_dropout=mlp_dropout,
                  attention_dropout=attention_dropout,
                  partition_size=partition_size,
                  input_grid_size=input_size,
                  n_layers=num_layers,
                  p_stochastic=p_stochastic[p_idx : p_idx + num_layers],
              ),
          )
          input_size = self.blocks[-1].grid_size
          p_idx += num_layers
  
      # see https://github.com/google-research/maxvit/blob/da76cf0d8a6ec668cc31b399c4126186da7da944/maxvit/models/maxvit.py#L1137-L1158
      # for why there is Linear -> Tanh -> Linear
      self.classifier = nn.Sequential(
          nn.AdaptiveAvgPool2d(1),
          nn.Flatten(),
          nn.LayerNorm(block_channels[-1]),
          nn.Linear(block_channels[-1], block_channels[-1]),
          nn.Tanh(),
          nn.Linear(block_channels[-1], num_classes, bias=False),
      )
  
      self._init_weights()
  
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 MaxVit object from inherited base classes: nn with example: obj = MaxVit(input_size, stem_channels, partition_size, block_channels, block_layers, head_dim, stochastic_depth_prob, norm_layer, activation_layer, squeeze_ratio, expansion_ratio, mlp_ratio, mlp_dropout, attention_dropout, num_classes)