code: def forward(self, x):
      b, c, hh, ww, u, h = *x.shape, self.u, self.heads
  
      q = self.to_q(x)
      k = self.to_k(x)
      v = self.to_v(x)
  
      q = self.norm_q(q)
      v = self.norm_v(v)
  
      q = rearrange(q, "b (h k) hh ww -> b h k (hh ww)", h=h)
      k = rearrange(k, "b (u k) hh ww -> b u k (hh ww)", u=u)
      v = rearrange(v, "b (u v) hh ww -> b u v (hh ww)", u=u)
  
      k = k.softmax(dim=-1)
  
      λc = einsum("b u k m, b u v m -> b k v", k, v)
      Yc = einsum("b h k n, b k v -> b h v n", q, λc)
  
      if self.local_contexts:
          v = rearrange(v, "b u v (hh ww) -> b u v hh ww", hh=hh, ww=ww)
          λp = self.pos_conv(v)
          Yp = einsum("b h k n, b k v n -> b h v n", q, λp.flatten(3))
      else:
          n, m = self.rel_pos.unbind(dim=-1)
          rel_pos_emb = self.rel_pos_emb[n, m]
          λp = einsum("n m k u, b u v m -> b n k v", rel_pos_emb, v)
          Yp = einsum("b h k n, b n k v -> b h v n", q, λp)
  
      Y = Yc + Yp
      out = rearrange(Y, "b h v (hh ww) -> b (h v) hh ww", hh=hh, ww=ww)
      return out
  
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 2 for calling the LambdaLayer obj's underlying member method code with example usage: obj.forward(x) and returns: out