Quoting from the einsum
documentation page:
To enable and control broadcasting, use an ellipsis. Default NumPy-style broadcasting is done by adding an ellipsis to the left of each term, like np.einsum('...ii->...i', a)
. To take the trace along the first and last axes, you can do np.einsum('i...i', a)
, or to do a matrix-matrix product with the left-most indices instead of rightmost, one can do np.einsum('ij...,jk...->ik...', a, b)
.
Later, an example is given:
>>> a = np.arange(25).reshape(5,5)
>>> np.einsum('...j->...', a)
array([ 10, 35, 60, 85, 110])
The equivalent MATLAB code to this example would be:
>> a = reshape(0:24, [5,5]).';
>> sum(a,2).'
ans =
10 35 60 85 110
Several things to note:
- The ellipsis operator (
...
) should not be understood as "range", but as "whatever needs to be there".
- "Broadcasting" refers to automatic replication of an array along the relevant dimension so that the mathematical operation is defined. This is a feature that exists in MATLAB since R2016b (called "implicit expansion").
- You may notice several transpose operations (
.'
) in the MATLAB equivalent. This is because numpy arrays are row-major while MATLAB array are column-major. Practically, while the underlying data has the same sequential order, a numpy array appears transposed compared to MATLAB. The transpositions were done so that arrays appear the same in intermediate stages.
Another example from these docs is:
>>> a = np.arange(6).reshape((3,2))
>>> b = np.arange(12).reshape((4,3))
>>> np.einsum('ki,jk->ij', a, b)
array([[10, 28, 46, 64],
[13, 40, 67, 94]])
>>> np.einsum('ki,...k->i...', a, b)
array([[10, 28, 46, 64],
[13, 40, 67, 94]])
>>> np.einsum('k...,jk', a, b)
array([[10, 28, 46, 64],
[13, 40, 67, 94]])
Which can be written as follows in MATLAB:
A = reshape(0:5, [2 3]).';
B = reshape(0:11, [3 4]).';
A.' * B.'
permute(sum( permute(A, [3 1 2]) .* B,2), [3 1 2])
shiftdim(sum( shiftdim(A, -1) .* B, 2), 2)
Several things to note:
- When going from
np.einsum('ki,jk->ij', a, b)
to np.einsum('ki,...k->i...', a, b)
you can see that the j
-th dimension is replaced with ...
. The fact that both these examples have a ->
in them, means it's in explicit mode.
- When going from
np.einsum('ki,jk->ij', a, b)
to np.einsum('k...,jk', a, b)
, you can see that now the i
-th dimension is replaced with ...
. The omission of ->...j
simply demonstrates implicit mode (where the output dimensions are ordered alphabetically).