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 forward(self, x: Tensor) -> Tensor:
      x = self.stem(x)
      for block in self.blocks:
          x = block(x)
      x = self.classifier(x)
      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 MaxVit obj's underlying member method code with example usage: obj.forward(x) and returns: x