code: def __init__(
      self,
      *,
      image_size,
      patch_size,
      num_classes,
      dim,
      heads,
      num_hierarchies,
      block_repeats,
      mlp_mult=4,
      channels=3,
      dim_head=64,
      dropout=0.0,
  ):
      super().__init__()
      assert (
          image_size % patch_size
      ) == 0, "Image dimensions must be divisible by the patch size."
      num_patches = (image_size // patch_size) ** 2
      patch_dim = channels * patch_size**2
      fmap_size = image_size // patch_size
      blocks = 2 ** (num_hierarchies - 1)
  
      seq_len = (
          fmap_size // blocks
      ) ** 2  # sequence length is held constant across hierarchy
      hierarchies = list(reversed(range(num_hierarchies)))
      mults = [2**i for i in reversed(hierarchies)]
  
      layer_heads = list(map(lambda t: t * heads, mults))
      layer_dims = list(map(lambda t: t * dim, mults))
      last_dim = layer_dims[-1]
  
      layer_dims = [*layer_dims, layer_dims[-1]]
      dim_pairs = zip(layer_dims[:-1], layer_dims[1:])
  
      self.to_patch_embedding = nn.Sequential(
          Rearrange(
              "b c (h p1) (w p2) -> b (p1 p2 c) h w",
              p1=patch_size,
              p2=patch_size,
          ),
          LayerNorm(patch_dim),
          nn.Conv2d(patch_dim, layer_dims[0], 1),
          LayerNorm(layer_dims[0]),
      )
  
      block_repeats = cast_tuple(block_repeats, num_hierarchies)
  
      self.layers = nn.ModuleList([])
  
      for level, heads, (dim_in, dim_out), block_repeat in zip(
          hierarchies, layer_heads, dim_pairs, block_repeats
      ):
          is_last = level == 0
          depth = block_repeat
  
          self.layers.append(
              nn.ModuleList(
                  [
                      Transformer(
                          dim_in,
                          seq_len,
                          depth,
                          heads,
                          mlp_mult,
                          dropout,
                      ),
                      Aggregate(dim_in, dim_out)
                      if not is_last
                      else nn.Identity(),
                  ]
              )
          )
  
      self.mlp_head = nn.Sequential(
          LayerNorm(last_dim),
          Reduce("b c h w -> b c", "mean"),
          nn.Linear(last_dim, num_classes),
      )
  
imports: from functools import partial
  import torch
  from torch import nn, einsum
  from einops import rearrange
  from einops.layers.torch import Rearrange, Reduce
label_desc: Use image node_id 1 to create a new NesT object from inherited base classes: nn with example: obj = NesT()