rolfmorel wrote: Fixed now. Will explain tomorrow why this works / it has to be this way. Same fix will need to happen for `linalg.batch_matmul` which I will do before this PR is ready to merge. https://github.com/llvm/llvm-project/pull/126377