Как я и обещал, готова аннотированная наивная имплементация Mamba, на пальцах.
Здесь две тетрадки, одна про теорию, вторая - с разобранной реализацией мамбы и тренировкой на MNIST: предсказываем вторую половину картинки по первой.
В качестве baseline - трансформер-декодер, который почему то не шибко смог, возможно баг. Я не стал разбираться, т.к. и так затянул с разбором. Зато мамба смогла молодцом. Замечания и багфиксы приветствуются. Репосты тоже :)
Рад если окажется полезным!

На первой картинке верхний ряд - ground truth, нижний - mamba prediction по первой половине пикселей (верхней части цифры).
@toshoseti