import torch
from nwt_pytorch import Memcodes
codebook = Memcodes(
    dim = 512,            # dimension of incoming features (codebook dimension will be dim / heads)
    heads = 8,            # head dimension, which is equivalent ot number of codebooks
    num_codes = 1024,     # number of codes per codebook
    temperature = 1.      # gumbel softmax temperature
)
x = torch.randn(1, 1024, 512)
out, codebook_indices = codebook(x) # (1, 1024, 512), (1, 1024, 8)
# (batch, seq, dimension), (batch, seq, heads)
# reconstruct output from codebook indices (codebook indices are autoregressed out from an attention net in paper)
assert torch.allclose(codebook.get_codes_from_indices(codebook_indices), out)