It's not perfect, but if you define a signature like `(m,n)->k` and call the kernel passing in an `out` array, then the value of `k` passed to the kernel will come from the shape of the `out` array. So you could have the shape calculation logic in the wrapper, create the output array there, and always call the kernel with an `out` parameter.