About the linalg module#
All the wrappers in this module (with the notable exception of
short and simple wrappers. Most are <5 lines and focus on providing a
comfortable and accessible API, not on extending the capabilities of the
The skeleton of such wrappers follows this pattern:
If no dims are provided, try to get default values for them. This is the main “extension” provided by this module. Doing
inv(da, dims=["dim2", "dim2"])still feels like too much unnecessary information is needed. We believe that in the majority of cases the dimensions corresponding to the matrices can be guessed with simple logic. Therefore, as explained in Default dimensions section we provide a way for users to define functions that set or infer the default matrix dimension names. We expect this to be particularly helpful in workflows that chain multiple operations on the same matrix and its results.
If required, set the dimensions of the output, generating new dimension names from the provided ones.
einsum_path are the only ones that don’t follow this pattern, and
they diverge significantly from it as we believe is required.
xarray already has an
xarray.dot, but it has two important
limitations. It reduces all the dimensions provided (so it can’t be used to
implement a matrix multiplication for example) and it takes a single list of dimensions,
common in for all inputs (you can’t reduce the “dim1” dimension for the 1st input but not the 2nd).
This means that a lot of the capabilities of
einsum are already covered by
by the functions in the
einops module. However, we believe that the cases that
are not covered (such as matrix multiplication or outer product)
are relevant enough to write this
It’s logic is similar to that of the einops module but with different
syntax, now adapted to
einsum instead of
einops. Here multiple inputs are supported
but decomposition/stacking is not, so dimensions being a list of lists means multiple outputs,
Moreover, most of the logic is abstracted into
which enables us to support both
einsum_path with very little extra code for
each. We hope that having
einsum_path will be helpful and accelerate operations
for our users.
As we explained in the introduction, we believe that allowing users to define the logic used to choose the dimensions on which to operate by default can be very powerful, even a game changer.
After considering some alternatives, we ended up deciding on providing the infrastructure and encouraging monkeypatching as a way for users to customize this logic. We want the functions in this module to have a simple call signature as close as possible to their numpy counterpart. Using alternatives like inheritance would require restructuring the module into a class and class methods. More functional approaches require adding an argument to the functions, if we ask users to pass an extra argument then it can be directly the dimensions on which to operate.