zhangyiss/mlx-examples
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx-examples.git synced 2025-08-13 12:46:43 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

nits

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

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

@ -65,6 +65,7 @@ def silu(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)
C = mx.expand_dims(C, axis=2) C = mx.expand_dims(C, axis=2)