code: def __init__(
      self,
      dim,
      *,
      dim_k,
      n=None,
      r=None,
      heads=4,
      dim_out=None,
      dim_u=1,
  ):
      super().__init__()
      dim_out = default(dim_out, dim)
      self.u = dim_u  # intra-depth dimension
      self.heads = heads
  
      assert (
          dim_out % heads
      ) == 0, "values dimension must be divisible by number of heads for multi-head query"
      dim_v = dim_out // heads
  
      self.to_q = nn.Conv2d(dim, dim_k * heads, 1, bias=False)
      self.to_k = nn.Conv2d(dim, dim_k * dim_u, 1, bias=False)
      self.to_v = nn.Conv2d(dim, dim_v * dim_u, 1, bias=False)
  
      self.norm_q = nn.BatchNorm2d(dim_k * heads)
      self.norm_v = nn.BatchNorm2d(dim_v * dim_u)
  
      self.local_contexts = exists(r)
      if exists(r):
          assert (r % 2) == 1, "Receptive kernel size should be odd"
          self.pos_conv = nn.Conv3d(
              dim_u, dim_k, (1, r, r), padding=(0, r // 2, r // 2)
          )
      else:
          assert exists(n), "You must specify the window size (n=h=w)"
          rel_lengths = 2 * n - 1
          self.rel_pos_emb = nn.Parameter(
              torch.randn(rel_lengths, rel_lengths, dim_k, dim_u)
          )
          self.rel_pos = calc_rel_pos(n)
  
file_path: lambda-networks/lambda_networks/lambda_networks.py
imports: import torch
  from torch import nn, einsum
  from einops import rearrange
label_desc: Use image node_id 1 to create a new LambdaLayer object from inherited base classes: nn with example: obj = LambdaLayer(dim)