[kingstnap]

Just for receipts I span up a basic test for XLA:

Can you figure out that 16 matrix vector multiplications into a concatenate is the same as concatenating first into a matrix matrix operation which can go on the GEMM. This is like the most basic thing you can imagine doing.

Turns out no it doesn't and there is a 23% performance difference by moving the concatenate up in the python code in my specific test.

Not to say that it didn't recognize it. I had an agent look at the XLA and the graph actually does a partial fusion into a sum and stack. But does not realize the whole thing is just a GEMM.

In more complicated examples the differences you can get can be much larger.