{"library":"product-key-memory","title":"Product Key Memory","description":"A PyTorch implementation of Product Key Memory (PKM), an external memory module for neural networks with fast nearest-neighbor lookup via product quantization. Current version 0.3.0, requires Python >=3.6. Active development by lucidrains.","language":"python","status":"active","last_verified":"Sat May 09","install":{"commands":["pip install product-key-memory"],"cli":null},"imports":["from product_key_memory import ProductKeyMemory"],"auth":{"required":false,"env_vars":[]},"quickstart":{"code":"import torch\nfrom product_key_memory import ProductKeyMemory\n\nmodel = ProductKeyMemory(\n    dim=512,\n    num_keys=512,\n    topk=32,\n    dim_head=64,\n    heads=8,\n    use_layernorm=True\n)\nx = torch.randn(1, 128, 512)\noutput, loss = model(x, x)\nprint(output.shape)  # (1, 128, 512)\nprint(loss.shape)    # (1,)","lang":"python","description":"Creates a PKM module, applies it to random input, returns output and auxiliary loss.","tag":null,"tag_description":null,"last_tested":null,"results":[]},"compatibility":null}