RFC: consider adding the specialized methods for missing matrix products

[NeilGirdhar]

Background

The matmul function. combines seven cases of matrix products in one function. These seven cases can be cleanly organized as examples of four possibly-vectorized operations:

• vector-dot, (does not support vectorization, but vector_dot covers this perfectly)
• matrix-matrix product (covers vectorization competely),
• matrix-vector-product (only supports vectorization in the matrix term), and
• vector-matrix-product (only supports vectorization in the matrix term).

Unfortunately, the decomposition isn't clean and the last two have operations have missing vectorization cases. For example,

def matrix_vector_mul(x: RealArray, y: RealArray) -> RealArray:
    """Return the matrix-vector product.

    This is xp.einsum("...ij,...j->...i", x, y).
    """
    xp = get_namespace(x, y)
    y = xp.reshape(y, (*y.shape[:-1], 1, y.shape[-1]))
    return xp.sum(x * y, axis=-1)

These would subsume cases 3 and 4.

Proposal

Consider adding matrix_vector_mul and vector_matrix_mul?

[rgommers]

It looks like this overlaps with matvec and vecmat, which were just added in NumPy 2.2.0: numpy/numpy#25675. Can you please review that PR and the mailing list discussion linked in the PR description, and confirm that that is what you're proposing here @NeilGirdhar?

[NeilGirdhar]

That's exactly what I want, yes. Will these be added to the Array API? Also, will matmul get the complex conjugation behavior?

[rgommers]

That's exactly what I want, yes. Will these be added to the Array API?

Great, thanks for confirming. I think what is needed first is to add those functions to other libraries. We can't standardize something that is only present in NumPy. Usually other libraries are happy to match NumPy if the API looks clean and there are no blockers like it not working well for accelerators or JITs (which I wouldn't expect to be an issue in this case). So it's "only" a matter of asking JAX, PyTorch et al. if these functions can be added and then someone implementing them.

[jakevdp]

JAX has already added the new functions in jax-ml/jax#25390

[rgommers]

We discussed this in today's community meeting. The feeling was that, beyond this not being actionable if the functions in question aren't yet present in most array libraries, it may be better to focus on einsum for the standard rather than a larger selection of particular functions like matvec.

Hence I propose we close this issue - I hope that's fine with you @NeilGirdhar? In case a function gets universally implemented across array libraries and it has become popular to use, we can always reconsider in the future.

[NeilGirdhar]

Yes, that works too! Hope you end up with a nice interface for einsum!

[rgommers]

Great, thanks Neil!