MzeroMiko / mamba-mini

An efficient pytorch implementation of selective scan in one file, works with both cpu and gpu, with corresponding mathematical derivation. It is probably the code which is the most close to selective_scan_cuda in mamba.
71 stars 1 forks source link

Is it possible for initialized the hidden space with given state (hprefix in your example)? #5

Closed pUmpKin-Co closed 7 months ago

pUmpKin-Co commented 7 months ago

Thanks for your work! Help me a alot. It may just need to replace the following line with given state I think?

hprefix = us.new_zeros((B, G, D, N), dtype=torch.float)
MzeroMiko commented 7 months ago

yes, you are right