How to implement custom operations with different input and output tensor shape? E.g., Matrix Multiplication

[JasperJin01]

Hello,

I was reviewing the test-customop.c example file and noticed that the ggml_map_custom1-3 functions seem to require the destination tensor (dst) to have the same shape as the input tensors. However, in my use case, I need the input tensors and the output tensor to have different shapes.

Specifically, my use case can be simplified to a "custom matrix multiplication problem." I would like to perform matrix multiplication using two input tensors a and b as the matrices, with the output tensor dst being the product of these matrices. The tensors a, b, and dst have different shapes. How can this be achieved? (P.S. I know that ggml includes the matrix multiplication like function ggml_mul_mat, but I indeed need to implement a custom version of matrix multiplication.)

Thank you for your assistance!