Skip to content

Preferred Operations

Leah Weisburn edited this page Dec 10, 2024 · 3 revisions

Multiplying matrices

Two arguments

Due to readability the binary @ operator is preferred over numpy.dot or numpy.matmul, i.e. please write:

A @ B

instead of

numpy.dot(A, B)
numpy.matmul(A, B)

Note that these are the equivalent when multiplying 2 dimensional matrices. In higher dimensions, these functions are different: An example adapted from here, rewriting the functions as einsums, is helpful to understand the differences between each:

numpy.allclose(numpy.einsum('ijk,ijk->ijk', a,b), a*b)                   # True 
numpy.allclose(numpy.einsum('ijk,ikl->ijl', a,b), a@b)                   # True
numpy.allclose(numpy.einsum('ijk,ikl->ijl', a,b), numpy.matmul(a,b))     # True
numpy.allclose(numpy.einsum('ijk,lkm->ijlm',a,b), a.dot(b))              # True

Note that numpy.matmul and @ are equivalent in higher dimensions, while numpy.dot is different.

Multiple arguments

In the case of writing a matrix product over multiple arguments of unknown size numpy.linalg.multi_dot is preferred since it reorders the evaluation of arguments to optimize performance. The matrix product is associative, but depending on the order of evaluation the amount of necessary multiplications and summations drastically differs; e.g. A @ (B @ v) is faster than (A @ B) @ v, where A, B are matrices and v is a vector.

See here the documentation of multi_dot

multi_dot chains numpy.dot and uses optimal parenthesization of the matrices [1] [2]. Depending on the shapes of the matrices, this can speed up the multiplication a lot.

This means if the shapes are unknown then prefer

multi_dot(A, B, C)

over

A @ B @ C

especially if performance matters.

Clone this wiki locally