gates (A matrix) with a shape of batch * dim * dim * seqlen

[WeihanLikk]

Thank you for your excellent work!

I was wondering if it’s possible to modify your code to handle a state-space model case where the gates (A matrix) have a more general shape of batch × dim × dim × seqlen?