Paper - Pytorch
Project description
MambaFormer
Implementation of MambaFormer in Pytorch ++ Zeta from the paper: "Can Mamba Learn How to Learn? A Comparative Study on In-Context Learning Tasks"
install
pip3 install mamba-former
usage
import torch
from mamba_former.main import MambaFormer
# Forward pass example
x = torch.randint(1, 1000, (1, 100)) # Token
# Tokens are integrers
# Model
model = MambaFormer(
dim = 512,
num_tokens = 1000,
depth = 6,
d_state = 512,
d_conv = 128,
heads = 8,
dim_head = 64,
return_tokens = True
)
# Forward
out = model(x)
print(out)
print(out.shape)
License
MIT
Project details
Release history Release notifications | RSS feed
Download files
Download the file for your platform. If you're not sure which to choose, learn more about installing packages.
Source Distribution
mamba_former-0.0.1.tar.gz
(3.6 kB
view hashes)
Built Distribution
Close
Hashes for mamba_former-0.0.1-py3-none-any.whl
Algorithm | Hash digest | |
---|---|---|
SHA256 | 530d0e61010d75800616291ea0f8fd8f279affea84d450cf753047c9a2a02ae6 |
|
MD5 | b00f2ceb1cc87e165d82c52cdec682fd |
|
BLAKE2b-256 | db88e5bfe08d11fbe71345d4d7794b4226034a4b592335bef7d68fc2a7591fcd |