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