code: def __init__(
      self,
      *,
      num_classes,
      dim,
      depth,
      heads,
      reduction_factor,
      window_size=None,
      iwsa_dim_key=32,
      iwsa_dim_value=32,
      ssa_dim_key=32,
      ssa_dim_value=32,
      ff_expansion_factor=4,
      channels=3,
      dropout=0.0,
  ):
      super().__init__()
      self.to_patches = nn.Conv2d(channels, dim, 7, stride=4, padding=3)
  
      assert isinstance(
          depth, tuple
      ), "depth needs to be tuple if integers indicating number of transformer blocks at that stage"
  
      num_stages = len(depth)
      dims = tuple(map(lambda i: (2**i) * dim, range(num_stages)))
  
      hyperparams_per_stage = [
          heads,
          ssa_dim_key,
          ssa_dim_value,
          reduction_factor,
          iwsa_dim_key,
          iwsa_dim_value,
          window_size,
      ]
  
      hyperparams_per_stage = list(
          map(
              partial(cast_tuple, length=num_stages),
              hyperparams_per_stage,
          )
      )
      assert all(
          tuple(
              map(
                  lambda arr: len(arr) == num_stages,
                  hyperparams_per_stage,
              )
          )
      )
  
      self.layers = nn.ModuleList([])
  
      for ind, (
          layer_dim,
          layer_depth,
          layer_heads,
          layer_ssa_dim_key,
          layer_ssa_dim_value,
          layer_ssa_reduction_factor,
          layer_iwsa_dim_key,
          layer_iwsa_dim_value,
          layer_window_size,
      ) in enumerate(zip(dims, depth, *hyperparams_per_stage)):
          is_last = ind == (num_stages - 1)
  
          self.layers.append(
              nn.ModuleList(
                  [
                      Transformer(
                          dim=layer_dim,
                          depth=layer_depth,
                          heads=layer_heads,
                          ff_expansion_factor=ff_expansion_factor,
                          dropout=dropout,
                          ssa_dim_key=layer_ssa_dim_key,
                          ssa_dim_value=layer_ssa_dim_value,
                          ssa_reduction_factor=layer_ssa_reduction_factor,
                          iwsa_dim_key=layer_iwsa_dim_key,
                          iwsa_dim_value=layer_iwsa_dim_value,
                          iwsa_window_size=layer_window_size,
                          norm_output=not is_last,
                      ),
                      Downsample(layer_dim, layer_dim * 2)
                      if not is_last
                      else None,
                  ]
              )
          )
  
      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
  from einops import rearrange, repeat
  from einops.layers.torch import Rearrange, Reduce
label_desc: Use image node_id 1 to create a new ScalableViT object from inherited base classes: nn with example: obj = ScalableViT()