[mcemilg]

Looks wonderful. But I would like to add this, I hate einops, it doesn't make it simple to read unfortunately.

[andy99]

I re-implemented Mamba myself and this was the first time I had ever worked with einops/einsum. I'm 50/50 on them after this. I found them relatively easy to look at and understand the intent (possibly more so than other representations), but talking extra time to transforms into other primitives (loops, multiplication, etc). I belive using torch.einsum is generally well optimized as well compared to naively looping. All said, I don't know if I'd use it myself working from scratch but it's interesting to know and if I was working in python I might try comparing the speed of einops/sum vs other ways.

[sjkoelle]

disagree