File size: 863 Bytes
d8c733f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 | from Vathos.blocks import *
try:
from mamba_ssm import Mamba
except ImportError:
raise ImportError("If you're trying to use Mamba, please install mamba-ssm, via pip install mamba-ssm. \n"
"probably you'll need to also install the right torch version:\n"
"pip install -q -U torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu126")
class MambaMixer(Layer):
def __init__(self, d_model, expand=4, d_state=64, k=3):
super(MambaMixer).__init__()
self.expand = expand
self.d_model = d_model
self.d_state = d_state
self.k = k
self.mamba = Mamba(
d_model=d_model,
d_state=d_state,
d_conv=d_state,
expand=expand,
)
def forward(self, x):
x = self.mamba(x)
return x
|