jtuyls wrote: > Thanks, it would be great to support `collapse_shape` as well. Sure, I can implement `collapse_shape` as a follow-up. Could you help merge this already in the meantime? https://github.com/llvm/llvm-project/pull/164438