{"id":28063,"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.","status":"active","version":"0.3.0","language":"python","source_language":"en","source_url":"https://github.com/lucidrains/product-key-memory","tags":["pytorch","memory","neural-network","product-quantization","lucidrains"],"install":[{"cmd":"pip install product-key-memory","lang":"bash","label":"Default install"}],"dependencies":[{"reason":"Core dependency; PKM module uses PyTorch tensors and operations.","package":"torch","optional":false}],"imports":[{"note":"","wrong":"","symbol":"ProductKeyMemory","correct":"from product_key_memory import ProductKeyMemory"}],"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."},"warnings":[{"fix":"Update code to expect only two return values. If you used aux_loss, it is now included in loss.","message":"In v0.3.0, the forward method signature changed: previously returned (output, loss, aux_loss), now returns (output, loss). The aux_loss is no longer returned separately.","severity":"breaking","affected_versions":"<0.3.0"},{"fix":"Ensure input is 3D: unsqueeze if necessary.","message":"The module expects input shape (batch, seq_len, dim). Common mistake: passing (batch, dim) for single timestep leads to errors.","severity":"gotcha","affected_versions":"all"},{"fix":"Use heads=1 or remove argument; check documentation.","message":"The argument `heads` may be deprecated in future versions; the library is moving to a more efficient single-head implementation.","severity":"deprecated","affected_versions":">=0.2.10"}],"env_vars":null,"last_verified":"2026-05-09T00:00:00.000Z","next_check":"2026-08-07T00:00:00.000Z","problems":[{"fix":"Upgrade to latest: pip install --upgrade product-key-memory","cause":"Installed an older version (pre-0.2.0) that used a different module name or structure.","error":"ImportError: cannot import name 'ProductKeyMemory' from 'product_key_memory'"},{"fix":"Ensure both are on same device: model = model.to('cuda'); x = x.to('cuda')","cause":"Model and input tensors on different devices (CPU vs GPU).","error":"RuntimeError: Expected tensor to be on the same device, but found at least two devices"},{"fix":"Remove return_loss; loss is now always computed and returned as second element.","cause":"Using deprecated keyword from older version (v0.1.x).","error":"TypeError: forward() got an unexpected keyword argument 'return_loss'"}],"ecosystem":"pypi","meta_description":null,"install_score":null,"install_tag":null,"quickstart_score":null,"quickstart_tag":null}