code: def __init__(
      self,
      *,
      num_classes,
      dim,
      depth,
      dim_head=32,
      dim_conv_stem=None,
      window_size=7,
      mbconv_expansion_rate=4,
      mbconv_shrinkage_rate=0.25,
      dropout=0.1,
      channels=3,
      num_register_tokens=4,
  ):
      super().__init__()
      assert isinstance(
          depth, tuple
      ), "depth needs to be tuple if integers indicating number of transformer blocks at that stage"
      assert num_register_tokens > 0
  
      # convolutional stem
  
      dim_conv_stem = default(dim_conv_stem, dim)
  
      self.conv_stem = Sequential(
          nn.Conv2d(channels, dim_conv_stem, 3, stride=2, padding=1),
          nn.Conv2d(dim_conv_stem, dim_conv_stem, 3, padding=1),
      )
  
      # variables
  
      num_stages = len(depth)
  
      dims = tuple(map(lambda i: (2**i) * dim, range(num_stages)))
      dims = (dim_conv_stem, *dims)
      dim_pairs = tuple(zip(dims[:-1], dims[1:]))
  
      self.layers = nn.ModuleList([])
  
      # window size
  
      self.window_size = window_size
  
      self.register_tokens = nn.ParameterList([])
  
      # iterate through stages
  
      for ind, ((layer_dim_in, layer_dim), layer_depth) in enumerate(
          zip(dim_pairs, depth)
      ):
          for stage_ind in range(layer_depth):
              is_first = stage_ind == 0
              stage_dim_in = layer_dim_in if is_first else layer_dim
  
              conv = MBConv(
                  stage_dim_in,
                  layer_dim,
                  downsample=is_first,
                  expansion_rate=mbconv_expansion_rate,
                  shrinkage_rate=mbconv_shrinkage_rate,
              )
  
              block_attn = Attention(
                  dim=layer_dim,
                  dim_head=dim_head,
                  dropout=dropout,
                  window_size=window_size,
                  num_registers=num_register_tokens,
              )
              block_ff = FeedForward(dim=layer_dim, dropout=dropout)
  
              grid_attn = Attention(
                  dim=layer_dim,
                  dim_head=dim_head,
                  dropout=dropout,
                  window_size=window_size,
                  num_registers=num_register_tokens,
              )
              grid_ff = FeedForward(dim=layer_dim, dropout=dropout)
  
              register_tokens = nn.Parameter(
                  torch.randn(num_register_tokens, layer_dim)
              )
  
              self.layers.append(
                  ModuleList(
                      [
                          conv,
                          ModuleList([block_attn, block_ff]),
                          ModuleList([grid_attn, grid_ff]),
                      ]
                  )
              )
  
              self.register_tokens.append(register_tokens)
  
      # mlp head out
  
      self.mlp_head = nn.Sequential(
          Reduce("b d h w -> b d", "mean"),
          nn.LayerNorm(dims[-1]),
          nn.Linear(dims[-1], num_classes),
      )
  
imports: from functools import partial
  import torch
  from torch import nn, einsum
  import torch.nn.functional
  from torch.nn import Module, ModuleList, Sequential
  from einops import rearrange, repeat, reduce, pack, unpack
  from einops.layers.torch import Rearrange, Reduce
label_desc: Use image node_id 1 to create a new MaxViT object from inherited base classes: Module with example: obj = MaxViT()