class_docstr: MobileViT.
  Paper: https://arxiv.org/abs/2110.02178
  Based on: https://github.com/chinhsuanwu/mobilevit-pytorch
code: def __init__(
      self,
      image_size,
      dims,
      channels,
      num_classes,
      expansion=4,
      kernel_size=3,
      patch_size=(2, 2),
      depths=(2, 4, 3),
  ):
      super().__init__()
      assert len(dims) == 3, "dims must be a tuple of 3"
      assert len(depths) == 3, "depths must be a tuple of 3"
  
      ih, iw = image_size
      ph, pw = patch_size
      assert ih % ph == 0 and iw % pw == 0
  
      init_dim, *_, last_dim = channels
  
      self.conv1 = conv_nxn_bn(3, init_dim, stride=2)
  
      self.stem = nn.ModuleList([])
      self.stem.append(MV2Block(channels[0], channels[1], 1, expansion))
      self.stem.append(MV2Block(channels[1], channels[2], 2, expansion))
      self.stem.append(MV2Block(channels[2], channels[3], 1, expansion))
      self.stem.append(MV2Block(channels[2], channels[3], 1, expansion))
  
      self.trunk = nn.ModuleList([])
      self.trunk.append(
          nn.ModuleList(
              [
                  MV2Block(channels[3], channels[4], 2, expansion),
                  MobileViTBlock(
                      dims[0],
                      depths[0],
                      channels[5],
                      kernel_size,
                      patch_size,
                      int(dims[0] * 2),
                  ),
              ]
          )
      )
  
      self.trunk.append(
          nn.ModuleList(
              [
                  MV2Block(channels[5], channels[6], 2, expansion),
                  MobileViTBlock(
                      dims[1],
                      depths[1],
                      channels[7],
                      kernel_size,
                      patch_size,
                      int(dims[1] * 4),
                  ),
              ]
          )
      )
  
      self.trunk.append(
          nn.ModuleList(
              [
                  MV2Block(channels[7], channels[8], 2, expansion),
                  MobileViTBlock(
                      dims[2],
                      depths[2],
                      channels[9],
                      kernel_size,
                      patch_size,
                      int(dims[2] * 4),
                  ),
              ]
          )
      )
  
      self.to_logits = nn.Sequential(
          conv_1x1_bn(channels[-2], last_dim),
          Reduce("b c h w -> b c", "mean"),
          nn.Linear(channels[-1], num_classes, bias=False),
      )
  
imports: import torch
  import torch.nn
  from einops import rearrange
  from einops.layers.torch import Reduce
label_desc: Use image node_id 1 to create a new MobileViT object from inherited base classes: nn with example: obj = MobileViT(image_size, dims, channels, num_classes, expansion, kernel_size, patch_size, depths)