RFC: add `kron` to compute the Kronecker product

[lucascolley]

Prior art

• NumPy - https://numpy.org/doc/stable/reference/generated/numpy.kron.html
• CuPy - https://docs.cupy.dev/en/latest/reference/generated/cupy.kron.html
• Dask - dask/dask#3657
• JAX - https://jax.readthedocs.io/en/latest/_autosummary/jax.numpy.kron.html
• PyTorch - https://pytorch.org/docs/stable/generated/torch.kron.html
• Tensorflow - https://github.com/tensorflow/tensorflow/blob/v2.16.1/tensorflow/python/ops/numpy_ops/np_math_ops.py#L436-L469

Motivation

We have an awkward function kron in scipy.linalg. We would like to deprecate the function and direct users to np.kron instead, like we have done for np.triu and np.tril. However, unlike those functions, xp.kron is not in the standard, hence we have been hesitant to deprecate in case it does not make it into the standard¹. Experiments done by @ilayn in scipy/scipy#20077 demonstrate that there are big performance gains to be had with a native array library implementation. Also, an array-agnostic implementation by @mdhaber in scipy/scipy#20077 (comment) demonstrates that at least a Python-level solution is easily implementable in standard array libraries.

Whether kron can be included in the standard or not, it would be very useful for us to know either way, so that we can plan whether we should deprecate scipy.linalg.kron or not.

Footnotes

1. this is a concern because scipy.linalg is a feasible home for an array-agnostic kron if it isn't in the standard - more info at https://github.com/scipy/scipy/issues/20077#issuecomment-1949281961 ↩

[asmeurer]

Two questions here:

1. Where is kron used, both in scipy and other libraries?
2. Can kron be implemented in terms of existing array API functions?

[mdhaber]

I can only answer the second question: yes, e.g. scipy/scipy#20077 (comment). There may be a performance benefit to library-specific implementation, though.

[lucascolley]

To partly answer (1), in SciPy production code:

• scipy.linalg.kron is used in:
  • scipy.linalg.solve_discrete_lyapunov
• scipy.sparse.kron is used in:
  • scipy.sparse.linalg.LaplacianNd
  • scipy.sparse.kronsum
• numpy.kron is unused

All 3 appear in either tests, benchmarks or tutorials too.

[asmeurer]

There were some discussions previously about making an array-api-helpers package that contains various useful functions that can be implemented on top of the array API.

Also, I don't know if this has ever been discussed, but NumPy itself could implement the array API for certain functions that are implemented in pure Python, similar to SciPy (np.kron is implemented in pure Python). Then libraries like scipy building on top of the array API could just call np.kron on the array API inputs.

Of course, that isn't to necessarily say it shouldn't be standardized. It appears to be implemented in many packages. But I think there's also a question of how often it is used. I'm +0 to adding it to the standard.

[lucascolley]

For the record I'm also +0 for the same reasons.

@ilayn can speak to the viewpoint that np.kron should be rewritten in a compiled language for performance and thus that it belongs in the standard.

There were some discussions previously about making an array-api-helpers package that contains various useful functions that can be implemented on top of the array API.

Yeah, I've talked about "array-api-extra" in the linked SciPy issue. Supposing such a package is created, I think the sensible solution here would be to add Matt's array-agnostic implementation of kron there. Then we can safely deprecate scipy.linalg.kron and perhaps kron is moved into the standard at a later point.

Also, I don't know if this has ever been discussed, but NumPy itself could implement the array API for certain functions that are implemented in pure Python, similar to SciPy (np.kron is implemented in pure Python). Then libraries like scipy building on top of the array API could just call np.kron on the array API inputs.

Cool idea, although this feels like quite a departure from the mental model I've been working with. For someone coming from another array library ecosystem, if kron isn't in the standard, there should be an array-agnostic way of calling kron that doesn't depend on NumPy. So I don't think this feels like the way forward to me.

[ilayn]

Independent from the array api discussion the footnote is a big no for me.

this is a concern because scipy.linalg is a feasible home for an array-agnostic kron if it isn't in the standard

This is not correct, it is an historical artifact from times where numpy and scipy distinction was forming. Also this is not a concern for me as a scipy maintainer. We never decided on making scipy as a home for feasibility. Moreover, I am very strong -1 on making scipy as a casual drop-in basket for things that are not in array API.

This particular point needs to be made crystal clear somewhere. Because we are discussing this all the time and it is draining morale for no good reason. I don't have any opinion whether it should go into standard or not. But it does not belong in SciPy since it is an array generator. And array consumers should not attempt doing this with low performant ways.

if kron isn't in the standard, there should be an array-agnostic way of calling kron that doesn't depend on NumPy.

This does not imply it should be in SciPy. That's point of disagreement. SciPy is not a complement to array API. It is an array consumer. So architecturally this is already an anti-pattern.

[lucascolley]

My justification for saying "feasible" in the footnote is available to read in scipy/scipy#20077, so I would not like to go over that again. Just to clear up a factual point here:

This does not imply it should be in SciPy. That's point of disagreement.

This is not true, I agree that kron does not belong in SciPy, which is why I have put effort into exploring the alternatives of adding it to the standard or this "array-api-extra" idea. It only reaches the "feasible" bar for me due to the circumstances that kron currently exists within SciPy, and SciPy is making lots of its algorithms array-agnostic.

[ilayn]

My justification for saying "feasible" in the footnote is available to read in scipy/scipy#20077,

That's exactly is the disagreement. I need to emphasize, I don't randomly go around deprecate things in scipy. So in the meantime, I am trying to explain to you now many many times that it is not something I'm just feeling like deprecating it. I need you to also understand why I am taking that position. You also need to grant me some credit that I do know why I am proposing to deprecate it.

I still need an answer why kron is a matter of array API that you have so far not answered. That's why I conceded from that discussion because it is revolving. So far it seems like you want the whole catalogue to be supported by array API and that part is kind of getting tiring now because that is not going to happen as I explained many times. Scipy.linalg api is not ready to be frozen or take all of array api. We are still working on batched input and many other details. We need to get rid of parts that does not fit. We cannot have this discussion everytime. So you have to leave some room for not being able to implement array API. Otherwise I don't need to be in this issue.

[izaid]

Just popping by to say I think the "array-api-extra" / "array-api-helpers" package would be very valuable! So many things can already be implemented with what exists, and there's a good argument that functions should just go live there unless it's either impossible or too inefficient to implement them that way.

I've used kron. It's valuable. I'd have to double-check if it can be implemented just with array operations (probably can), and if it can I think it should be.

[lucascolley]

If it helps, I would be happy to start up this array-api-extra package and help maintain it going forward. I assume the consortium would rather it was something official that lived under the org, though.

[rgommers]

Now that array-api-extra is alive (🎉) and kron is in the process of being implemented there, I'll go ahead and close this feature request. We discussed it today in a consortium call, and the consensus was that it doesn't meet the criteria for inclusion, since it's (a) not used much, (b) can be implemented in terms of other functions already in the standard, and (c) is significantly more niche than the other constructor functions present in the standard.

Thanks for the discussion everyone!