[Feature] Kronecker Product

[cvnad1]

@awni @angeloskath

I do not see the Kron function implemented in MLX, but I feel that it would be a great addition (Given its significance in science and math). If it is not yet implemented, I would love to add it.

[angeloskath]

Are you thinking of something more than (a[:, None, :, None] * b[None, :, None, :]).reshape(M, N)?

Personally, I think broadcasting is more powerful and more explicit so I am not sure. Wdyt? As an example, I never remember the extra rules these functions have so I have to search the docs. Can I perform broadcasting kronecker? What if I want to do Kronecker of batched vectors MxA kron MxB -> Mx(A*B) ?

[cvnad1]

@angeloskath Broadcasting is a good way and in fact largely similar to the official JAX implementation of the Kronecker product. Providing Kronecker product as an in-built function is still beneficial given its presence and use in various fields.

I did not come across a batched version of Kron yet. If it is fine to add the Kron function to the library, I can implement the non-batched version similar to what you are thinking and refer to the JAX library whether they have implemented anything related. wdyt?

[cvnad1]

@angeloskath May I know what you think of the above?

[angeloskath]

Personally I am overly conservative in these things. ie I tend to want to keep the ops as few as possible. I think our policy so far has been that if it is a numpy op then we will add it. In this case it is implemented by numpy and torch with exactly the same API so feel free to go ahead and add it.

[cvnad1]

@angeloskath thanks for the green signal !!! I too align with you. I am also raising issues with adding functions or features that are either implemented in Torch or Numpy. Keeping MLX minimal but effective is always the priority.

Kronecker is one such essential function and would love to add this. Again, Thanks for always supporting me with your valuable guidance and suggestions. I am actively contributing to MLX because of your support.