code: def __init__(
      self,
      dim,
      depth,
      heads=8,
      ff_expansion_factor=4,
      dropout=0.0,
      ssa_dim_key=32,
      ssa_dim_value=32,
      ssa_reduction_factor=1,
      iwsa_dim_key=32,
      iwsa_dim_value=32,
      iwsa_window_size=None,
      norm_output=True,
  ):
      super().__init__()
      self.layers = nn.ModuleList([])
      for ind in range(depth):
          is_first = ind == 0
  
          self.layers.append(
              nn.ModuleList(
                  [
                      ScalableSelfAttention(
                          dim,
                          heads=heads,
                          dim_key=ssa_dim_key,
                          dim_value=ssa_dim_value,
                          reduction_factor=ssa_reduction_factor,
                          dropout=dropout,
                      ),
                      FeedForward(
                          dim,
                          expansion_factor=ff_expansion_factor,
                          dropout=dropout,
                      ),
                      PEG(dim) if is_first else None,
                      FeedForward(
                          dim,
                          expansion_factor=ff_expansion_factor,
                          dropout=dropout,
                      ),
                      InteractiveWindowedSelfAttention(
                          dim,
                          heads=heads,
                          dim_key=iwsa_dim_key,
                          dim_value=iwsa_dim_value,
                          window_size=iwsa_window_size,
                          dropout=dropout,
                      ),
                  ]
              )
          )
  
      self.norm = ChanLayerNorm(dim) if norm_output else nn.Identity()
  
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 Transformer object from inherited base classes: nn with example: obj = Transformer(dim, depth, heads, ff_expansion_factor, dropout, ssa_dim_key, ssa_dim_value, ssa_reduction_factor, iwsa_dim_key, iwsa_dim_value, iwsa_window_size, norm_output)