zhangyiss/mlx-examples
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx-examples.git synced 2025-06-27 19:31:20 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

nits

Browse Source
This commit is contained in:
Goekdeniz-Guelmez 2024-12-10 18:15:12 +01:00
parent 9f8a6a3509
commit b10afe3662
1 changed files with 4 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
llms/mlx_lm/models/mamba2.py
View File

@ -62,6 +62,7 @@ class MambaRMSNormGated(nn.Module):
def silu(x): def silu(x):
return x * mx.sigmoid(x) return x * mx.sigmoid(x)
def ssd(x, A, B, C, chunk_size): def ssd(x, A, B, C, chunk_size):
batch, seqlen, nheads, dim = x.shape batch, seqlen, nheads, dim = x.shape
B = mx.expand_dims(B, axis=2) B = mx.expand_dims(B, axis=2)
@ -175,7 +176,7 @@ class Mamba2Block(nn.Module):
# Calculate split indices and slice tensors # Calculate split indices and slice tensors
z = proj[..., :self.d_inner] z = proj[..., :self.d_inner]
x_conv = proj[..., self.d_inner:self.d_inner + (self.d_inner + 2 * self.d_state)] x_conv = proj[..., self.d_inner:self.d_inner + (self.d_inner + 2 * self.n_groups * self.d_state)]
dt = proj[..., -self.n_heads:] dt = proj[..., -self.n_heads:]
# Process time steps # Process time steps