[ndarray] General matmul broadcasting bug about burn

[lancelet]

OK; PR raised here: #1679

I'm completely new to burn, so please let me know what extra tests, examples, docs, etc. should be added.

I saw a mention in another PR of improving the docs for matmul. Should I put that into this PR as well?

[lancelet]

EDIT: I have a WIP branch here: https://github.com/lancelet/burn/tree/matmul-broadcasting
In that branch, I use strides for the non-matrix dimensions (ie. the batch, depth, channel, or whatever) to handle arbitrary broadcasting.

It's getting the correct result for the above example, but it's failing on some autoregressive test cases with an error about an incompatible memory layout. This occurs when I try to flatten / reshape the arrays, in the same way that the original code does. Still investigating.

[laggui]

Hi @lancelet, thanks for looking into this! 🙂 I'll take a quick look at your PR for this issue specifically.

We do want to improve our documentation regarding the general broadcasting semantics in burn, and also matmul. If you want to take a stab at it you could open another PR!