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

optimizing the code for faster inference but still generates giberish

Browse Source
This commit is contained in:
Goekdeniz-Guelmez 2024-12-12 21:08:33 +01:00
parent c1d9ec329c
commit a883e39f41
1 changed files with 2 additions and 14 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

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

@ -109,23 +109,11 @@ class DepthWiseConv1d(nn.Module):
else: else:
x = mx.pad(x, [(0, 0), (K - 1, 0), (0, 0)]) x = mx.pad(x, [(0, 0), (K - 1, 0), (0, 0)])
# Adjust the weight tensor to match the input channels y = mx.conv_general(x, self.weight, groups=C)
if C != self.channels: y = y + self.bias
adjusted_weight = self.weight[:C, :, :]
else:
adjusted_weight = self.weight
y = mx.conv_general(x, adjusted_weight, groups=C)
if self.bias is not None:
# Adjust the bias to match the input channels
adjusted_bias = self.bias[:C] if C != self.channels else self.bias
y = y + adjusted_bias
return y, x[:, -K + 1:, :] return y, x[:, -K + 1:, :]
class Mamba2Block(nn.Module): class Mamba2Block(nn.Module):
def __init__(self, args: ModelArgs): def __init__(self, args: ModelArgs):
super().__init__() super().__init__()