import torch
from lambda_networks import LambdaLayer
layer = LambdaLayer(
    dim = 32,       # channels going in
    dim_out = 32,   # channels out
    n = 64,         # size of the receptive window - max(height, width)
    dim_k = 16,     # key dimension
    heads = 4,      # number of heads, for multi-query
    dim_u = 1       # 'intra-depth' dimension
)
x = torch.randn(1, 32, 64, 64)
layer(x) # (1, 32, 64, 64)

import torch
from lambda_networks import LambdaLayer
layer = LambdaLayer(
    dim = 32,
    dim_out = 32,
    r = 23,         # the receptive field for relative positional encoding (23 x 23)
    dim_k = 16,
    heads = 4,
    dim_u = 4
)
x = torch.randn(1, 32, 64, 64)
layer(x) # (1, 32, 64, 64)

from lambda_networks import λLayer

import tensorflow as tf
from lambda_networks.tfkeras import LambdaLayer
layer = LambdaLayer(
    dim_out = 32,
    r = 23,
    dim_k = 16,
    heads = 4,
    dim_u = 1
)
x = tf.random.normal((1, 64, 64, 16)) # channel last format
layer(x) # (1, 64, 64, 32)