code: def __init__(
      self,
      *,
      depth,
      sm_dim,
      lg_dim,
      sm_enc_params,
      lg_enc_params,
      cross_attn_heads,
      cross_attn_depth,
      cross_attn_dim_head=64,
      dropout=0.0,
  ):
      super().__init__()
      self.layers = nn.ModuleList([])
      for _ in range(depth):
          self.layers.append(
              nn.ModuleList(
                  [
                      Transformer(
                          dim=sm_dim, dropout=dropout, **sm_enc_params
                      ),
                      Transformer(
                          dim=lg_dim, dropout=dropout, **lg_enc_params
                      ),
                      CrossTransformer(
                          sm_dim=sm_dim,
                          lg_dim=lg_dim,
                          depth=cross_attn_depth,
                          heads=cross_attn_heads,
                          dim_head=cross_attn_dim_head,
                          dropout=dropout,
                      ),
                  ]
              )
          )
  
imports: import torch
  from torch import nn, einsum
  import torch.nn.functional
  from einops import rearrange, repeat
  from einops.layers.torch import Rearrange
label_desc: Use image node_id 1 to create a new MultiScaleEncoder object from inherited base classes: nn with example: obj = MultiScaleEncoder()