issue_comments: 821744766

Would there be any interest in an "xarray linalg library"? Some kind of library of that sort would be really useful to ArviZ and our users, I'll probably write some very basic functions myself on a minimal library or add them to ArviZ directly, but if there are other people interested we can try to find synergies and get something more general.

In our case, 99% of the time we want to "batch" over the chain and draw dimensions. So things as simple as:

def xarray_dot(a, b, dim=None): if dim is None: a_dims = set(a.dims) b_dims = set(b.dims) dim_set = a_dims.intersection(b_dims) - {"chain", "draw"} if len(dim_set) == 1: dim = dim_set.pop() else: ValueError if not isinstance(dim, str): raise ValueError return xr.apply_ufunc( np.einsum, "...i,...i", a, b, input_core_dims=[[], [dim], [dim]] )

are already extremely useful.

Being able to invert, cholesky... without using apply_ufunc would be convenient, even more so if we define some conventions. i.e. I generally use dim and dim bis for covariance matrices so they have shape chain, draw, ..., dim, dim bis, so I can invert those independently of their number and order of dimensions without even needing to say which are the "matrix" dimensions.

In our case, being able to ds.transpose("dim1", "dim2") and have it ignore chain, draw would also be convenient, which is probably a very ArviZ specific situation, but maybe there are other libraries/people that commonly have some batch dimensions, even if they use different names for them or if the user chooses those names freely.