code: def __init__(
      self,
      *,
      dim,
      depth=2,
      dim_context=None,
      num_latents=32,
      dim_head=64,
      heads=8,
      ff_mult=4,
      use_flash_attn=False,
  ):
      super().__init__()
      dim_context = default(dim_context, dim)
  
      self.proj_context = (
          nn.Linear(dim_context, dim)
          if dim_context != dim
          else nn.Identity()
      )
  
      self.latents = nn.Parameter(torch.randn(num_latents, dim))
      nn.init.normal_(self.latents, std=0.02)
  
      self.layers = nn.ModuleList([])
      for _ in range(depth):
          self.layers.append(
              nn.ModuleList(
                  [
                      Attention(
                          dim=dim,
                          dim_head=dim_head,
                          heads=heads,
                          use_flash=use_flash_attn,
                          cross_attn_include_queries=True,
                      ),
                      FeedForward(dim=dim, mult=ff_mult),
                  ]
              )
          )
  
      self.norm = RMSNorm(dim)
  
imports: from collections import namedtuple
  from functools import wraps
  import torch
  import torch.nn.functional
  from einops import rearrange, repeat
  from einops.layers.torch import Rearrange
  from packaging import version
  from torch import einsum, nn
label_desc: Use image node_id 1 to create a new PerceiverResampler object from inherited base classes: nn with example: obj = PerceiverResampler()