[BUG] Why is the backpropagation calculation so slow when I use the mamba network?

[1325116124]

When I used the mamba network, I defined a loss to test backpropagation and found that the calculation was very slow. Setting the len length to 1024 requires a long waiting time. code show as below：

`import torch
import torch.nn as nn
from zeta.nn import MambaBlock

block = MambaBlock(dim=512, depth=1)
x = torch.randn(1, 1024, 512)
target = torch.randn(1, 1024, 512)
loss_fn = nn.MSELoss()

y = block(x)
loss = loss_fn(y, target)
loss.backward()
print("Output shape:", y.shape)
print("Loss value:", loss.item())
`

Upvote & Fund

• We're using Polar.sh so you can upvote and help fund this issue.
• We receive the funding once the issue is completed & confirmed by you.
• Thank you in advance for helping prioritize & fund our backlog.

[github-actions]

Hello there, thank you for opening an Issue ! 🙏🏻 The team was notified and they will get back to you asap.

[kyegomez]

@1325116124 its using the mamba scan, or SSM, it should be updated soon!

[Alex-Naxitus]

I had the same problem. I don't suppose you managed to fix it right ?

[kyegomez]

@Alex-Naxitus yes, it should be updated now!

[kyegomez]

@Alex-Naxitus @1325116124 it was the backscan that was really slow, let me know so I can close this issuee!

[github-actions]

Stale issue message