Update mamba.py

mamba.py

@@ -160,7 +160,7 @@ class MambaBlock(nn.Module):
 
         deltaBC = self.x_proj(x) 
 
-        delta, B, C = torch.split(deltaBC, [self.config.dt_rank, self.config.d_state, self.config.d_state], dim=-1) # (B, L, dt_rank), (B, L, N), (B, L, N)
+        delta, B, C = torch.split(deltaBC, [self.config.dt_rank, self.config.d_state, self.config.d_state], dim=-1) 
         delta = F.softplus(self.dt_proj(delta)) 
 
         if self.config.pscan:
@@ -196,7 +196,7 @@ class MambaBlock(nn.Module):
 
         BX = deltaB * (x.unsqueeze(-1)) 
 
-        h = torch.zeros(x.size(0), self.config.d_inner, self.config.d_state, device=deltaA.device) # (B, ED, N)
+        h = torch.zeros(x.size(0), self.config.d_inner, self.config.d_state, device=deltaA.device) 
         hs = []
 
         for t in range(0, L):
@@ -233,10 +233,10 @@ class MambaBlock(nn.Module):
         z = F.silu(z)
 
         output = y * z
-        output = self.out_proj(output) # (B, D)
+        output = self.out_proj(output) 
 
         
-        inputs = torch.cat([inputs[:, :, 1:], x_cache], dim=2) # (B, ED, d_conv-1)
+        inputs = torch.cat([inputs[:, :, 1:], x_cache], dim=2) 
         cache = (h, inputs)
         
         return output, cache