Search code examples
pythonnumpynumpy-einsumnumpy-ufunc

What is the precise meaning of question marks in an einsum-like gufunc signature?


For example:

np.arange(3)@np.arange(2)
# Traceback (most recent call last):
#  File "<stdin>", line 1, in <module>
# ValueError: matmul: Input operand 1 has a mismatch in its core dimension 0, with gufunc signature (n?,k),(k,m?)->(n?,m?) (size 2 is different from 3)

Please note: I know what's wrong and what the error message means. I'm just interested in the precise syntax. What do the question marks signify?


Solution

  • These indicate dimensions that are allowed to be missing. If such a dimension is missing, it will be treated as if it were present with length 1, but the corresponding dimension will be removed from the output.

    In the matmul signature (n?,k),(k,m?)->(n?,m?), the first dimension of the first argument and the second dimension of the second argument are both allowed to be missing. If they are, the first argument will be treated as a row vector, and the second argument will be treated as a column vector.

    Quoting NEP 20 — Expansion of Generalized Universal Function Signatures,

    1. Possibly missing dimensions. This part is almost entirely driven by the wish to wrap matmul in a gufunc. matmul stands for matrix multiplication, and if it did only that, it could be covered with the signature (m,n),(n,p)->(m,p). However, it has special cases for when a dimension is missing, allowing either argument to be treated as a single vector, with the function thus becoming, effectively, vector-matrix, matrix-vector, or vector-vector multiplication (but with no broadcasting). To support this, it is suggested to allow postfixing a dimension name with a question mark to indicate that the dimension does not necessarily have to be present.

      With this addition, the signature for matmul can be expressed as (m?,n),(n,p?)->(m?,p?). This indicates that if, e.g., the second operand has only one dimension, for the purposes of the elementary function it will be treated as if that input has core shape (n, 1), and the output has the corresponding core shape of (m, 1). The actual output array, however, has the flexible dimension removed, i.e., it will have shape (..., m). Similarly, if both arguments have only a single dimension, the inputs will be presented as having shapes (1, n) and (n, 1) to the elementary function, and the output as (1, 1), while the actual output array returned will have shape (). In this way, the signature allows one to use a single elementary function for four related but different signatures, (m,n),(n,p)->(m,p), (n),(n,p)->(p), (m,n),(n)->(m) and (n),(n)->().