Add mlx mamba training script

alxndrTL / mamba.py

A simple and efficient Mamba implementation in pure PyTorch and MLX.

MIT License

960 stars 86 forks source link

Add mlx mamba training script #36

Closed michal1000w closed 4 months ago

michal1000w commented 4 months ago

Hi. Based on my previous PR with mamba LM pytorch training script I've created a similar mlx version. From what I've tested it works correctly and the code is very similar for better consistency.

alxndrTL commented 4 months ago

Hello, sorry for the late response, thank for the PR!