einsum for xarray

[fujiisoup]

Code Sample, a copy-pastable example if possible

I sometimes want to make more flexible dot product of two data arrays,
where we sum up along a part of common dimensions.

# Your code here
da_vals = np.arange(6 * 5 * 4).reshape((6, 5, 4))
da = DataArray(da_vals, dims=['x', 'y', 'z'])
dm_vals = np.arange(6 * 4).reshape((6, 4))
dm = DataArray(dm_vals, dims=['x', 'z'])

# I want something like this
da.dot(dm, 'z')  # -> dimensions of the output array: ['x', 'y']

It's an intermediate path of np.matmul and np.tensordot.

Is this feature sufficiently universal?

EDIT:
I just noticed dask does not have einsum yet (dask/dask#732).
We maybe need to wait or decide to support only numpy arrays.

[shoyer]

I like the idea of an einsum wrapper for xarray. This would probably be straightforward to build on top of the apply_ufunc machinery. I wouldn't worry about the lack of einsum for dask yet -- the linked dask issue already has several implementations, so we could easily add that later if desired.

There are at least two versions of this:

• A dim argument for DataArray.dot that indicates which dimension(s) to sum over.
• A function for applying this to many arrays simultaneously. Maybe xarray.dot(da1, da2, da3, dims=['x', 'y'])? This would use einsum internally but we don't need to specify the input dimensions again because we already have them on the DataArray objects.