code: def __init__(
      self, *, dim_k, n=None, r=None, heads=4, dim_out=None, dim_u=1
  ):
      super(LambdaLayer, self).__init__()
  
      self.out_dim = dim_out
      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"
      self.dim_v = dim_out // heads
      self.dim_k = dim_k
      self.heads = heads
  
      self.to_q = Conv2D(self.dim_k * heads, 1, use_bias=False)
      self.to_k = Conv2D(self.dim_k * dim_u, 1, use_bias=False)
      self.to_v = Conv2D(self.dim_v * dim_u, 1, use_bias=False)
  
      self.norm_q = BatchNormalization()
      self.norm_v = BatchNormalization()
  
      self.local_contexts = exists(r)
      if exists(r):
          assert (r % 2) == 1, "Receptive kernel size should be odd"
          self.pos_conv = Conv3D(dim_k, (1, r, r), padding="same")
      else:
          assert exists(
              n
          ), "You must specify the window length (n = h = w)"
          rel_length = 2 * n - 1
          self.rel_pos_emb = self.add_weight(
              name="pos_emb",
              shape=(rel_length, rel_length, dim_k, dim_u),
              initializer=initializers.random_normal,
              trainable=True,
          )
          self.rel_pos = calc_rel_pos(n)
  
file_path: lambda-networks/lambda_networks/tfkeras.py
imports: import tensorflow
  from einops.layers.tensorflow import Rearrange
  from tensorflow.keras.layers import Conv2D, BatchNormalization, Conv3D, ZeroPadding3D, Softmax, Lambda, Add, Layer
  from tensorflow.keras import initializers
  from tensorflow import einsum, nn, meshgrid
label_desc: Use image node_id 1 to create a new LambdaLayer object from inherited base classes: Layer with example: obj = LambdaLayer()